[Core][Disk] High disk tier on Azure

sky/clouds/azure.py

@@ -361,7 +361,7 @@ def _failover_disk_tier() -> Optional[resources_utils.DiskTier]:
                 start_index += 1
             assert False, 'Low disk tier should always be supported on Azure.'
 
-        return {
+        resources_vars = {
             'instance_type': r.instance_type,
             'custom_resources': custom_resources,
             'num_gpus': acc_count,
@@ -377,6 +377,12 @@ def _failover_disk_tier() -> Optional[resources_utils.DiskTier]:
             'resource_group': f'{cluster_name.name_on_cloud}-{region_name}',
         }
 
+        disk_performance_tier = Azure._get_disk_performance_tier(r.disk_tier)
+        if disk_performance_tier is not None:
+            resources_vars['disk_performance_tier'] = disk_performance_tier
+
+        return resources_vars
+
     def _get_feasible_launchable_resources(
         self, resources: 'resources.Resources'
     ) -> 'resources_utils.FeasibleResources':
@@ -600,10 +606,10 @@ def check_disk_tier(
             disk_tier: Optional[resources_utils.DiskTier]) -> Tuple[bool, str]:
         if disk_tier is None or disk_tier == resources_utils.DiskTier.BEST:
             return True, ''
-        if disk_tier == resources_utils.DiskTier.HIGH or disk_tier == resources_utils.DiskTier.ULTRA:
+        if disk_tier == resources_utils.DiskTier.ULTRA:
             return False, (
-                'Azure disk_tier={high, ultra} is not supported now. '
-                'Please use disk_tier={low, medium, best} instead.')
+                'Azure disk_tier=ultra is not supported now. '
+                'Please use disk_tier={low, medium, high, best} instead.')
         # Only S-series supported premium ssd
         # see https://stackoverflow.com/questions/48590520/azure-requested-operation-cannot-be-performed-because-storage-account-type-pre  # pylint: disable=line-too-long
         if cls._get_disk_type(
@@ -631,8 +637,17 @@ def _get_disk_type(cls,
         # cannot be used as OS disks so we might need data disk support
         tier2name = {
             resources_utils.DiskTier.ULTRA: 'Disabled',
-            resources_utils.DiskTier.HIGH: 'Disabled',
+            resources_utils.DiskTier.HIGH: 'Premium_LRS',
             resources_utils.DiskTier.MEDIUM: 'Premium_LRS',
             resources_utils.DiskTier.LOW: 'Standard_LRS',
         }
         return tier2name[tier]
+
+    @classmethod
+    def _get_disk_performance_tier(
+            cls,
+            disk_tier: Optional[resources_utils.DiskTier]) -> Optional[str]:
+        tier = cls._translate_disk_tier(disk_tier)
+        if tier == resources_utils.DiskTier.HIGH:
+            return 'P50'
+        return None

sky/provision/azure/instance.py

@@ -18,6 +18,7 @@
 from sky.provision import common
 from sky.provision import constants
 from sky.utils import common_utils
+from sky.utils import subprocess_utils
 from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
@@ -274,6 +275,16 @@ def _create_instances(
         deployment_name=vm_name,
         parameters=parameters,
     ).wait()
+
+    if 'disk_performance_tier' in provider_config:
+        performance_tier = provider_config['disk_performance_tier']
+        disks = compute_client.disks.list_by_resource_group(resource_group)
+        for disk in disks:
+            name = disk.name
+            subprocess_utils.run_no_outputs(
+                f'az disk update -n {name} -g {resource_group} '
+                f'--set tier={performance_tier}')
+
     filters = {
         constants.TAG_RAY_CLUSTER_NAME: cluster_name_on_cloud,
         _TAG_SKYPILOT_VM_ID: vm_id

sky/templates/azure-ray.yml.j2

@@ -45,6 +45,9 @@ provider:
   # The upper-level SkyPilot code has make sure there will not be resource
   # leakage.
   disable_launch_config_check: true
+  {%- if disk_performance_tier is not none %}
+  disk_performance_tier: {{disk_performance_tier}}
+  {%- endif %}
 
 
 auth: