Enable EphemeralDiskPlacentType capability for the Azure VMs

lisa/features/disks.py

@@ -54,6 +54,9 @@ def get_all_disks(self) -> List[str]:
     def get_hardware_disk_controller_type(self) -> schema.DiskControllerType:
         raise NotImplementedError
 
+    def get_ephemeral_disk_placement_type(self) -> schema.EphemeralDiskPlacementType:
+        raise NotImplementedError
+
     def add_data_disk(
         self,
         count: int,
@@ -137,7 +140,8 @@ def get_os_disk_controller_type(self) -> schema.DiskControllerType:
 
 
 DiskEphemeral = partial(
-    schema.DiskOptionSettings, os_disk_type=schema.DiskType.Ephemeral
+    schema.DiskOptionSettings,
+    os_disk_type=schema.DiskType.Ephemeral,
 )
 DiskPremiumSSDLRS = partial(
     schema.DiskOptionSettings,

lisa/sut_orchestrator/azure/arm_template.bicep

@@ -98,7 +98,7 @@ func getEphemeralOSImage(node object) object => {
   name: '${node.name}-osDisk'
   diffDiskSettings: {
       option: 'local'
-      placement: 'CacheDisk'
+      placement: node.ephemeral_disk_placement_type
   }
   caching: 'ReadOnly'
   createOption: 'FromImage'

lisa/sut_orchestrator/azure/common.py

@@ -1068,6 +1068,7 @@ class AzureNodeArmParameter(AzureNodeSchema):
     os_disk_type: str = ""
     data_disk_type: str = ""
     disk_controller_type: str = ""
+    ephemeral_disk_placement_type: str = ""
     security_profile: Dict[str, Any] = field(default_factory=dict)
 
     @classmethod

lisa/sut_orchestrator/azure/features.py

@@ -1447,6 +1447,32 @@ def _call_requirement_method(
                 self.os_disk_size, capability.os_disk_size
             )
 
+        if value.os_disk_type == schema.DiskType.Ephemeral:
+            cap_ephemeral_disk_placement_type = capability.ephemeral_disk_placement_type
+            if isinstance(cap_ephemeral_disk_placement_type, search_space.SetSpace):
+                assert len(cap_ephemeral_disk_placement_type) > 0, (
+                    "capability should have at least one "
+                    "ephemeral disk placement type, but it's empty"
+                )
+            elif isinstance(
+                cap_ephemeral_disk_placement_type, schema.EphemeralDiskPlacementType):
+                cap_ephemeral_disk_placement_type = search_space.SetSpace[
+                    schema.EphemeralDiskPlacementType](
+                        is_allow_set=True, items=[cap_ephemeral_disk_placement_type])
+            else:
+                raise LisaException(
+                    "unknown ephemeral disk placement type "
+                    f"on capability, type: {cap_ephemeral_disk_placement_type}"
+                )
+
+            value.ephemeral_disk_placement_type = getattr(
+                search_space, f"{method.value}_setspace_by_priority"
+            )(
+                self.ephemeral_disk_placement_type,
+                capability.ephemeral_disk_placement_type,
+                schema.ephemeral_disk_placement_type_priority,
+            )
+
         value.data_disk_type = getattr(
             search_space, f"{method.value}_setspace_by_priority"
         )(self.data_disk_type, capability.data_disk_type, schema.disk_type_priority)
@@ -1665,6 +1691,11 @@ def get_hardware_disk_controller_type(self) -> Any:
         vm = get_vm(azure_platform, self._node)
         return vm.storage_profile.disk_controller_type
 
+    def get_ephemeral_disk_placement_type(self) -> Any:
+        azure_platform: AzurePlatform = self._platform  # type: ignore
+        vm = get_vm(azure_platform, self._node)
+        return vm.storage_profile.os_disk.diff_disk_settings.placement
+
     def _get_scsi_data_disks(self) -> List[str]:
         # This method restuns azure data disks attached to you given VM.
         # refer here to get data disks from folder /dev/disk/azure/scsi1

lisa/sut_orchestrator/azure/platform_.py

@@ -738,6 +738,19 @@ def _get_disk_controller_type(self, node: Node) -> str:
             node.log.debug(f"error on collecting disk controller type: {identifier}")
         return result
 
+    def _get_ephemeral_disk_placement_type(self, node: Node) -> str:
+        result: str = ""
+        try:
+            result = node.features[Disk].get_ephemeral_disk_placement_type()
+        except Exception as identifier:
+            # it happens on some error vms. Those error should be caught earlier in
+            # test cases not here. So ignore any error here to collect information only.
+            node.log.debug(
+                f"error on collecting ephemeral disk placement type:"
+                f" {identifier}"
+            )
+        return result
+
     def _get_kernel_version(self, node: Node) -> str:
         result: str = ""
 
@@ -1417,6 +1430,16 @@ def _create_node_arm_parameters(
         arm_parameters.os_disk_type = features.get_azure_disk_type(
             capability.disk.os_disk_type
         )
+        # Set Ephemeral Disk placement type
+        if arm_parameters.os_disk_type == schema.DiskType.Ephemeral:
+            assert isinstance(
+                capability.disk.ephemeral_disk_placement_type,
+                schema.EphemeralDiskPlacementType
+            )
+            arm_parameters.ephemeral_disk_placement_type = (
+                capability.disk.ephemeral_disk_placement_type.value
+            )
+
         assert isinstance(capability.disk.data_disk_type, schema.DiskType)
         arm_parameters.data_disk_type = features.get_azure_disk_type(
             capability.disk.data_disk_type
@@ -1794,14 +1817,40 @@ def _resource_sku_to_capability(  # noqa: C901
         else:
             node_space.disk.disk_controller_type.add(schema.DiskControllerType.SCSI)
 
+        # If EphemeralOSDisk is supported, then check for the placement type
         if azure_raw_capabilities.get("EphemeralOSDiskSupported", None) == "True":
-            # Check if CachedDiskBytes is greater than 30GB
-            # We use diff disk as cache disk for ephemeral OS disk
+            node_space.disk.os_disk_type.add(schema.DiskType.Ephemeral)
+            # Add the EphemeralDiskPlacementType
+            ephemeral_disk_placement_types = azure_raw_capabilities.get(
+                "SupportedEphemeralOSDiskPlacements", None)
+            if ephemeral_disk_placement_types:
+                for allowed_type in ephemeral_disk_placement_types.split(","):
+                    try:
+                        node_space.disk.ephemeral_disk_placement_type.add(
+                            schema.EphemeralDiskPlacementType(allowed_type)
+                        )
+                    except ValueError:
+                        self._log.error(
+                            f"'{allowed_type}' is not a known Ephemeral Disk Placement"
+                            f" Type "
+                            f"({[x for x in schema.EphemeralDiskPlacementType]})"
+                        )
+
+            # EphemeralDiskPlacementType can be - ResourceDisk, CacheDisk or NvmeDisk.
+            # Depending on that, "CachedDiskBytes" may or may not be found in
+            # capabilities.
+            # refer
+            # https://learn.microsoft.com/en-us/azure/virtual-machines/ephemeral-os-disks-faq
+            resource_disk_bytes = azure_raw_capabilities.get("MaxResourceVolumeMB", 0)
             cached_disk_bytes = azure_raw_capabilities.get("CachedDiskBytes", 0)
-            cached_disk_bytes_gb = int(int(cached_disk_bytes) / 1024 / 1024 / 1024)
-            if cached_disk_bytes_gb >= 30:
-                node_space.disk.os_disk_type.add(schema.DiskType.Ephemeral)
-                node_space.disk.os_disk_size = cached_disk_bytes_gb
+            nvme_disk_bytes = azure_raw_capabilities.get("NvmeDiskSizeInMiB", 0)
+            if nvme_disk_bytes:
+                node_space.disk.os_disk_size = int(int(nvme_disk_bytes) / 1024)
+            elif cached_disk_bytes:
+                node_space.disk.os_disk_size = int(
+                    int(cached_disk_bytes) / 1024 / 1024 / 1024)
+            else:
+                node_space.disk.os_disk_size = int(int(resource_disk_bytes) / 1024)
 
         # set AN
         if azure_raw_capabilities.get("AcceleratedNetworkingEnabled", None) == "True":
@@ -2047,6 +2096,15 @@ def _generate_max_capability(self, vm_size: str, location: str) -> AzureCapabili
         ](is_allow_set=True, items=[])
         node_space.disk.disk_controller_type.add(schema.DiskControllerType.SCSI)
         node_space.disk.disk_controller_type.add(schema.DiskControllerType.NVME)
+        node_space.disk.ephemeral_disk_placement_type = search_space.SetSpace[
+            schema.EphemeralDiskPlacementType
+        ](is_allow_set=True, items=[])
+        node_space.disk.ephemeral_disk_placement_type.add(
+            schema.EphemeralDiskPlacementType.Nvme)
+        node_space.disk.ephemeral_disk_placement_type.add(
+            schema.EphemeralDiskPlacementType.Cache)
+        node_space.disk.ephemeral_disk_placement_type.add(
+            schema.EphemeralDiskPlacementType.Resource)
         node_space.network_interface = schema.NetworkInterfaceOptionSettings()
         node_space.network_interface.data_path = search_space.SetSpace[
             schema.NetworkDataPath
@@ -2754,6 +2812,10 @@ def _set_disk_features(
             isinstance(node_space.disk.os_disk_type, search_space.SetSpace)
             and node_space.disk.os_disk_type.isunique(schema.DiskType.Ephemeral)
         ):
+            node_space.disk.ephemeral_disk_placement_type = search_space.SetSpace[
+                schema.EphemeralDiskPlacementType
+            ](is_allow_set=True, items=[node_space.disk.ephemeral_disk_placement_type])
+
             node_space.disk.os_disk_size = search_space.IntRange(
                 min=self._get_os_disk_size(azure_runbook)
             )