Azure: Add support for ARM64 Images

This is a draft PR, please don't merge it yet.

Initial attempt to add ARM64 support for Azure on CloudPub

Author: JAVGan

cloudpub/ms_azure/utils.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 import logging
 from operator import attrgetter
-from typing import Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 from deepdiff import DeepDiff
 
@@ -69,7 +69,22 @@ def __init__(
         super(AzurePublishingMetadata, self).__init__(**kwargs)
         self.__validate()
         # Adjust the x86_64 architecture string for Azure
-        self.architecture = "x64" if self.architecture == "x86_64" else self.architecture
+        arch = self.__convert_arch(self.architecture)
+        self.architecture = arch
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        if name == "architecture":
+            arch = self.__convert_arch(value)
+            value = arch
+        return super().__setattr__(name, value)
+
+    @staticmethod
+    def __convert_arch(arch: str) -> str:
+        converter = {
+            "x86_64": "x64",
+            "aarch64": "arm64",
+        }
+        return converter.get(arch, None) or arch
 
     def __validate(self):
         mandatory = [
@@ -91,9 +106,10 @@ def __validate(self):
 def get_image_type_mapping(architecture: str, generation: str) -> str:
     """Return the image type required by VMImageDefinition."""
     gen_map = {
-        "V1": f"{architecture}Gen1",
         "V2": f"{architecture}Gen2",
     }
+    if architecture == "x64":
+        gen_map.update({"V1": f"{architecture}Gen1"})
     return gen_map.get(generation, "")
 
 
@@ -185,6 +201,17 @@ def is_azure_job_not_complete(job_details: ConfigureStatus) -> bool:
     return False
 
 
+def is_legacy_gen_supported(metadata: AzurePublishingMetadata) -> bool:
+    """Return True when the legagy V1 SKU is supported, False otherwise.
+
+    Args:
+        metadata: The incoming publishing metadata.
+    Returns:
+        bool: True when V1 is supported, False otherwise.
+    """
+    return metadata.architecture == "x64" and metadata.support_legacy
+
+
 def prepare_vm_images(
     metadata: AzurePublishingMetadata,
     gen1: Optional[VMImageDefinition],
@@ -214,19 +241,21 @@ def prepare_vm_images(
         raise ValueError(msg)
 
     raw_source = source.to_json()
+    arch_name = metadata.architecture
+
     json_gen1 = {
-        "imageType": get_image_type_mapping(metadata.architecture, "V1"),
+        "imageType": get_image_type_mapping(arch_name, "V1"),
         "source": raw_source,
     }
     json_gen2 = {
-        "imageType": get_image_type_mapping(metadata.architecture, "V2"),
+        "imageType": get_image_type_mapping(arch_name, "V2"),
         "source": raw_source,
     }
 
     if metadata.generation == "V2":
         # In this case we need to set a V2 SAS URI
         gen2_new = VMImageDefinition.from_json(json_gen2)
-        if metadata.support_legacy:  # and in this case a V1 as well
+        if is_legacy_gen_supported(metadata):  # and in this case a V1 as well
             gen1_new = VMImageDefinition.from_json(json_gen1)
             return [gen2_new, gen1_new]
         return [gen2_new]
@@ -235,13 +264,25 @@ def prepare_vm_images(
         return [VMImageDefinition.from_json(json_gen1)]
 
 
+def _len_vm_images(disk_versions: List[DiskVersion]) -> int:
+    count = 0
+    for disk_version in disk_versions:
+        count = count + len(disk_version.vm_images)
+    return count
+
+
 def _build_skus(
     disk_versions: List[DiskVersion],
     default_gen: str,
     alt_gen: str,
     plan_name: str,
     security_type: Optional[List[str]] = None,
 ) -> List[VMISku]:
+    def get_skuid(arch):
+        if arch == "x64":
+            return plan_name
+        return f"{plan_name}-{arch.lower()}"
+
     sku_mapping: Dict[str, str] = {}
     # Update the SKUs for each image in DiskVersions if needed
     for disk_version in disk_versions:
@@ -254,10 +295,11 @@ def _build_skus(
             new_img_alt_type = get_image_type_mapping(arch, alt_gen)
 
             # we just want to add SKU whenever it's not set
+            skuid = get_skuid(arch)
             if vmid.image_type == new_img_type:
-                sku_mapping.setdefault(new_img_type, plan_name)
+                sku_mapping.setdefault(new_img_type, skuid)
             elif vmid.image_type == new_img_alt_type:
-                sku_mapping.setdefault(new_img_alt_type, f"{plan_name}-gen{alt_gen[1:]}")
+                sku_mapping.setdefault(new_img_alt_type, f"{skuid}-gen{alt_gen[1:]}")
 
     # Return the expected SKUs list
     res = [
@@ -290,16 +332,20 @@ def update_skus(
         The updated list with VMISkus.
     """
     if not old_skus:
+        # First image's gen
+        # generation = "V" + disk_versions[0].vm_images[0].image_type.split("Gen")[-1]
         alt_gen = "V2" if generation == "V1" else "V1"
         return _build_skus(
             disk_versions, default_gen=generation, alt_gen=alt_gen, plan_name=plan_name
         )
 
-    # If we have SKUs for both genenerations we don't need to update them as they're already
+    # Second case: All SKUs are provided
+    # If we have SKUs for each image we don't need to update them as they're already
     # properly set.
-    if len(old_skus) == 2:
+    if len(old_skus) == _len_vm_images(disk_versions):
         return old_skus
 
+    # Third case: Missing at least one SKU
     # Update SKUs to create the alternate gen.
     # The security type may exist only for Gen2, so it iterates over all gens to find it
     security_type = None
@@ -348,16 +394,17 @@ def create_disk_version_from_scratch(
         source (VMImageSource):
             The VMImageSource with the updated SAS URI.
     """
+    arch_name = metadata.architecture
     vm_images = [
         {
-            "imageType": get_image_type_mapping(metadata.architecture, metadata.generation),
+            "imageType": get_image_type_mapping(arch_name, metadata.generation),
             "source": source.to_json(),
         }
     ]
-    if metadata.support_legacy:
+    if is_legacy_gen_supported(metadata):
         vm_images.append(
             {
-                "imageType": get_image_type_mapping(metadata.architecture, "V1"),
+                "imageType": get_image_type_mapping(arch_name, "V1"),
                 "source": source.to_json(),
             }
         )
@@ -457,16 +504,17 @@ def create_vm_image_definitions(
     log.debug("Creating VMImageDefinitions for \"%s\"", metadata.destination)
     vm_images = []
 
+    arch_name = metadata.architecture
     vm_images.append(
         VMImageDefinition(
-            image_type=get_image_type_mapping(metadata.architecture, metadata.generation),
+            image_type=get_image_type_mapping(arch_name, metadata.generation),
             source=source.to_json(),
         )
     )
-    if metadata.support_legacy:  # Only True when metadata.generation == V2
+    if is_legacy_gen_supported(metadata):
         vm_images.append(
             VMImageDefinition(
-                image_type=get_image_type_mapping(metadata.architecture, "V1"),
+                image_type=get_image_type_mapping(arch_name, "V1"),
                 source=source.to_json(),
             )
         )

tests/ms_azure/conftest.py

@@ -333,6 +333,14 @@ def gen2_image(vmimage_source) -> Dict[str, Any]:
     }
 
 
+@pytest.fixture
+def arm_image(vmimage_source) -> Dict[str, Any]:
+    return {
+        "imageType": "arm64Gen2",
+        "source": vmimage_source,
+    }
+
+
 @pytest.fixture
 def disk_version(gen1_image: Dict[str, Any], gen2_image: Dict[str, Any]) -> Dict[str, Any]:
     return {
@@ -342,6 +350,15 @@ def disk_version(gen1_image: Dict[str, Any], gen2_image: Dict[str, Any]) -> Dict
     }
 
 
+@pytest.fixture
+def disk_version_arm64(arm_image):
+    return {
+        "versionNumber": "2.1.0",
+        "vmImages": [arm_image],
+        "lifecycleState": "generallyAvailable",
+    }
+
+
 @pytest.fixture
 def technical_config(disk_version: Dict[str, Any]) -> Dict[str, Any]:
     return {
@@ -540,11 +557,21 @@ def gen2_image_obj(gen2_image: Dict[str, Any]) -> VMImageDefinition:
     return VMImageDefinition.from_json(gen2_image)
 
 
+@pytest.fixture
+def arm_image_obj(arm_image: Dict[str, Any]) -> VMImageDefinition:
+    return VMImageDefinition.from_json(arm_image)
+
+
 @pytest.fixture
 def disk_version_obj(disk_version: Dict[str, Any]) -> DiskVersion:
     return DiskVersion.from_json(disk_version)
 
 
+@pytest.fixture
+def disk_version_arm64_obj(disk_version_arm64: Dict[str, Any]) -> DiskVersion:
+    return DiskVersion.from_json(disk_version_arm64)
+
+
 @pytest.fixture
 def vmimage_source_obj(vmimage_source: Dict[str, Any]) -> VMImageSource:
     return VMImageSource.from_json(vmimage_source)

tests/ms_azure/test_service.py

@@ -1327,7 +1327,7 @@ def test_is_submission_in_preview(
     @mock.patch("cloudpub.ms_azure.service.create_disk_version_from_scratch")
     @mock.patch("cloudpub.ms_azure.AzureService.filter_product_resources")
     @mock.patch("cloudpub.ms_azure.AzureService.get_product_plan_by_name")
-    def test_publish_live(
+    def test_publish_live_x64_only(
         self,
         mock_getprpl_name: mock.MagicMock,
         mock_filter: mock.MagicMock,
@@ -1405,3 +1405,93 @@ def test_publish_live(
         ]
         mock_submit.assert_has_calls(submit_calls)
         mock_ensure_publish.assert_called_once_with(product_obj.id)
+
+    @mock.patch("cloudpub.ms_azure.AzureService.ensure_can_publish")
+    @mock.patch("cloudpub.ms_azure.AzureService.get_submission_state")
+    @mock.patch("cloudpub.ms_azure.AzureService.diff_offer")
+    @mock.patch("cloudpub.ms_azure.AzureService.configure")
+    @mock.patch("cloudpub.ms_azure.AzureService.submit_to_status")
+    @mock.patch("cloudpub.ms_azure.utils.prepare_vm_images")
+    @mock.patch("cloudpub.ms_azure.service.is_sas_present")
+    @mock.patch("cloudpub.ms_azure.service.create_disk_version_from_scratch")
+    @mock.patch("cloudpub.ms_azure.AzureService.filter_product_resources")
+    @mock.patch("cloudpub.ms_azure.AzureService.get_product_plan_by_name")
+    def test_publish_live_arm64_only(
+        self,
+        mock_getprpl_name: mock.MagicMock,
+        mock_filter: mock.MagicMock,
+        mock_disk_scratch: mock.MagicMock,
+        mock_is_sas: mock.MagicMock,
+        mock_prep_img: mock.MagicMock,
+        mock_submit: mock.MagicMock,
+        mock_configure: mock.MagicMock,
+        mock_diff_offer: mock.MagicMock,
+        mock_getsubst: mock.MagicMock,
+        mock_ensure_publish: mock.MagicMock,
+        product_obj: Product,
+        plan_summary_obj: PlanSummary,
+        metadata_azure_obj: AzurePublishingMetadata,
+        technical_config_obj: VMIPlanTechConfig,
+        disk_version_arm64_obj: DiskVersion,
+        submission_obj: ProductSubmission,
+        azure_service: AzureService,
+    ) -> None:
+        metadata_azure_obj.overwrite = False
+        metadata_azure_obj.keepdraft = False
+        metadata_azure_obj.support_legacy = True
+        metadata_azure_obj.destination = "example-product/plan-1"
+        metadata_azure_obj.disk_version = "2.1.0"
+        metadata_azure_obj.architecture = "aarch64"
+        technical_config_obj.disk_versions = disk_version_arm64_obj
+        mock_getprpl_name.return_value = product_obj, plan_summary_obj
+        mock_filter.side_effect = [
+            [technical_config_obj],
+            [submission_obj],
+        ]
+        mock_getsubst.side_effect = ["preview", "live"]
+        mock_res_preview = mock.MagicMock()
+        mock_res_live = mock.MagicMock()
+        mock_res_preview.job_result = mock_res_live.job_result = "succeeded"
+        mock_submit.side_effect = [mock_res_preview, mock_res_live]
+        mock_is_sas.return_value = False
+        expected_source = VMImageSource(
+            source_type="sasUri",
+            os_disk=OSDiskURI(uri=metadata_azure_obj.image_path).to_json(),
+            data_disks=[],
+        )
+        disk_version_arm64_obj.vm_images[0] = VMImageDefinition(
+            image_type=get_image_type_mapping(metadata_azure_obj.architecture, "V2"),
+            source=expected_source.to_json(),
+        )
+        mock_prep_img.return_value = deepcopy(
+            disk_version_arm64_obj.vm_images
+        )  # During submit it will pop the disk_versions
+        technical_config_obj.disk_versions = [disk_version_arm64_obj]
+
+        # Test
+        azure_service.publish(metadata_azure_obj)
+        mock_getprpl_name.assert_called_once_with("example-product", "plan-1")
+        filter_calls = [
+            mock.call(product=product_obj, resource="virtual-machine-plan-technical-configuration"),
+            mock.call(product=product_obj, resource="submission"),
+        ]
+        mock_filter.assert_has_calls(filter_calls)
+        mock_is_sas.assert_called_once_with(
+            technical_config_obj,
+            metadata_azure_obj.image_path,
+        )
+        mock_prep_img.assert_called_once_with(
+            metadata=metadata_azure_obj,
+            gen1=None,
+            gen2=disk_version_arm64_obj.vm_images[0],
+            source=expected_source,
+        )
+        mock_disk_scratch.assert_not_called()
+        mock_diff_offer.assert_called_once_with(product_obj)
+        mock_configure.assert_called_once_with(resource=technical_config_obj)
+        submit_calls = [
+            mock.call(product_id=product_obj.id, status="preview"),
+            mock.call(product_id=product_obj.id, status="live"),
+        ]
+        mock_submit.assert_has_calls(submit_calls)
+        mock_ensure_publish.assert_called_once_with(product_obj.id)