Support placement type for disks

lisa/features/disks.py

@@ -137,7 +137,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/schema.py

@@ -446,11 +446,22 @@ class ResourceDiskType(str, Enum):
     NVME = constants.STORAGE_INTERFACE_TYPE_NVME
 
 
+class EphemeralDiskPlacementType(str, Enum):
+    Resource = constants.DISK_PLACEMENT_TYPE_RESOURCE
+    Cache = constants.DISK_PLACEMENT_TYPE_CACHE
+    Nvme = constants.DISK_PLACEMENT_TYPE_NVME
+
+
 disk_controller_type_priority: List[DiskControllerType] = [
     DiskControllerType.SCSI,
     DiskControllerType.NVME,
 ]
 
+ephemeral_disk_placement_type_priority: List[EphemeralDiskPlacementType] = [
+    EphemeralDiskPlacementType.Nvme,
+    EphemeralDiskPlacementType.Cache,
+    EphemeralDiskPlacementType.Resource,
+]
 
 os_disk_types: List[DiskType] = [
     DiskType.StandardHDDLRS,
@@ -559,6 +570,34 @@ class DiskOptionSettings(FeatureSettings):
             )
         ),
     )
+    ephemeral_disk_placement_type: Optional[
+        Union[
+            search_space.SetSpace[
+                EphemeralDiskPlacementType
+            ],
+            EphemeralDiskPlacementType,
+        ]
+    ] = field(  # type:ignore
+        default_factory=partial(
+            search_space.SetSpace,
+            items=[
+                EphemeralDiskPlacementType.Resource,
+                EphemeralDiskPlacementType.Cache,
+                EphemeralDiskPlacementType.Nvme,
+            ],
+        ),
+        metadata=field_metadata(
+            decoder=partial(
+                search_space.decode_nullable_set_space,
+                base_type=EphemeralDiskPlacementType,
+                default_values=[
+                    EphemeralDiskPlacementType.Resource,
+                    EphemeralDiskPlacementType.Cache,
+                    EphemeralDiskPlacementType.Nvme,
+                ],
+            )
+        ),
+    )
 
     def __eq__(self, o: object) -> bool:
         if not super().__eq__(o):
@@ -577,6 +616,7 @@ def __eq__(self, o: object) -> bool:
             and self.data_disk_size == o.data_disk_size
             and self.max_data_disk_count == o.max_data_disk_count
             and self.disk_controller_type == o.disk_controller_type
+            and self.ephemeral_disk_placement_type == o.ephemeral_disk_placement_type
         )
 
     def __repr__(self) -> str:
@@ -591,6 +631,7 @@ def __repr__(self) -> str:
             f"size: {self.data_disk_size},"
             f"max_data_disk_count: {self.max_data_disk_count},"
             f"disk_controller_type: {self.disk_controller_type}"
+            f"ephemeral_disk_placement_type: {self.ephemeral_disk_placement_type}"
         )
 
     def __str__(self) -> str:
@@ -641,6 +682,7 @@ def _get_key(self) -> str:
             f"{self.data_disk_caching_type}/{self.data_disk_iops}/"
             f"{self.data_disk_throughput}/{self.data_disk_size}/"
             f"{self.disk_controller_type}"
+            f"{self.ephemeral_disk_placement_type}"
         )
 
     def _call_requirement_method(
@@ -700,6 +742,17 @@ def _call_requirement_method(
                 capability.disk_controller_type,
                 disk_controller_type_priority,
             )
+        if (
+            self.ephemeral_disk_placement_type
+            or capability.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,
+                ephemeral_disk_placement_type_priority,
+            )
 
         return value
 

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'
@@ -333,7 +333,7 @@ resource nodes_data_disks 'Microsoft.Compute/disks@2022-03-02' = [
   }
 ]
 
-resource nodes_vms 'Microsoft.Compute/virtualMachines@2022-08-01' = [for i in range(0, node_count): {
+resource nodes_vms 'Microsoft.Compute/virtualMachines@2024-03-01' = [for i in range(0, node_count): {
   name: nodes[i].name
   location: nodes[i].location
   tags: combined_vm_tags

lisa/sut_orchestrator/azure/autogen_arm_template.json

@@ -5,8 +5,8 @@
   "metadata": {
     "_generator": {
       "name": "bicep",
-      "version": "0.30.23.60470",
-      "templateHash": "17909783643222378721"
+      "version": "0.32.4.45862",
+      "templateHash": "16398577375970436728"
     }
   },
   "functions": [
@@ -113,7 +113,7 @@
               "name": "[format('{0}-osDisk', parameters('node').name)]",
               "diffDiskSettings": {
                 "option": "local",
-                "placement": "CacheDisk"
+                "placement": "[parameters('node').ephemeral_disk_placement_type]"
               },
               "caching": "ReadOnly",
               "createOption": "FromImage",
@@ -685,7 +685,7 @@
         "count": "[length(range(0, variables('node_count')))]"
       },
       "type": "Microsoft.Compute/virtualMachines",
-      "apiVersion": "2022-08-01",
+      "apiVersion": "2024-03-01",
       "name": "[parameters('nodes')[range(0, variables('node_count'))[copyIndex()]].name]",
       "location": "[parameters('nodes')[range(0, variables('node_count'))[copyIndex()]].location]",
       "tags": "[variables('combined_vm_tags')]",
@@ -787,8 +787,8 @@
           "metadata": {
             "_generator": {
               "name": "bicep",
-              "version": "0.30.23.60470",
-              "templateHash": "12249187708601787514"
+              "version": "0.32.4.45862",
+              "templateHash": "7856159159103188049"
             }
           },
           "functions": [

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,36 @@ def _call_requirement_method(
                 self.os_disk_size, capability.os_disk_size
             )
 
+        # refer
+        # https://learn.microsoft.com/en-us/powershell/module/az.compute/set-azvmssstorageprofile?view=azps-13.0.0               # noqa: E501
+        # https://github.com/MicrosoftDocs/azure-compute-docs/blob/main/articles/virtual-machines/ephemeral-os-disks-faq.md      # noqa: E501
+        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 +1695,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,17 @@ 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"
+                           f" placement type: {identifier}")
+        return result
+
     def _get_kernel_version(self, node: Node) -> str:
         result: str = ""
 
@@ -1417,6 +1428,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 +1815,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 +2094,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 +2810,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)
             )

lisa/util/constants.py

@@ -183,3 +183,10 @@
 # StorageInterfaceTypes
 STORAGE_INTERFACE_TYPE_NVME = "NVMe"
 STORAGE_INTERFACE_TYPE_SCSI = "SCSI"
+
+# SupportedEphemeralOSDiskPlacements
+# refer
+# https://learn.microsoft.com/en-us/azure/virtual-machines/ephemeral-os-disks-faq
+DISK_PLACEMENT_TYPE_RESOURCE = "ResourceDisk"
+DISK_PLACEMENT_TYPE_CACHE = "CacheDisk"
+DISK_PLACEMENT_TYPE_NVME = "NvmeDisk"