skypilot-nightly 1.0.0.dev20241023__py3-none-any.whl → 1.0.0.dev20241025__py3-none-any.whl

sky/jobs/recovery_strategy.py

@@ -24,6 +24,7 @@ from sky.utils import common_utils
 from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
+    from sky import resources
     from sky import task as task_lib
 
 logger = sky_logging.init_logger(__name__)
@@ -327,8 +328,7 @@ class StrategyExecutor:
                         'Failure happened before provisioning. Failover '
                         f'reasons: {reasons_str}')
                     if raise_on_failure:
-                        raise exceptions.ProvisionPrechecksError(
-                            reasons=reasons)
+                        raise exceptions.ProvisionPrechecksError(reasons)
                     return None
                 logger.info('Failed to launch a cluster with error: '
                             f'{common_utils.format_exception(e)})')
@@ -382,7 +382,7 @@ class FailoverStrategyExecutor(StrategyExecutor, name='FAILOVER',
         # first retry in the same cloud/region. (Inside recover() we may not
         # rely on cluster handle, as it can be None if the cluster is
         # preempted.)
-        self._launched_resources: Optional['sky.resources.Resources'] = None
+        self._launched_resources: Optional['resources.Resources'] = None
 
     def _launch(self,
                 max_retry: Optional[int] = 3,

sky/provision/azure/azure-config-template.json

@@ -13,6 +13,12 @@
             "metadata": {
                 "description": "Subnet parameters."
             }
+        },
+        "nsgName": {
+            "type": "string",
+            "metadata": {
+                "description": "Name of the Network Security Group associated with the SkyPilot cluster."
+            }
         }
     },
     "variables": {
@@ -20,7 +26,7 @@
         "location": "[resourceGroup().location]",
         "msiName": "[concat('sky-', parameters('clusterId'), '-msi')]",
         "roleAssignmentName": "[concat('sky-', parameters('clusterId'), '-ra')]",
-        "nsgName": "[concat('sky-', parameters('clusterId'), '-nsg')]",
+        "nsgName": "[parameters('nsgName')]",
         "nsg": "[resourceId('Microsoft.Network/networkSecurityGroups', variables('nsgName'))]",
         "vnetName": "[concat('sky-', parameters('clusterId'), '-vnet')]",
         "subnetName": "[concat('sky-', parameters('clusterId'), '-subnet')]"

sky/provision/azure/config.py

@@ -8,7 +8,7 @@ import json
 from pathlib import Path
 import random
 import time
-from typing import Any, Callable
+from typing import Any, Callable, Tuple
 
 from sky import exceptions
 from sky import sky_logging
@@ -22,6 +22,7 @@ UNIQUE_ID_LEN = 4
 _DEPLOYMENT_NAME = 'skypilot-config'
 _LEGACY_DEPLOYMENT_NAME = 'ray-config'
 _RESOURCE_GROUP_WAIT_FOR_DELETION_TIMEOUT = 480  # 8 minutes
+_CLUSTER_ID = '{cluster_name_on_cloud}-{unique_id}'
 
 
 def get_azure_sdk_function(client: Any, function_name: str) -> Callable:
@@ -41,11 +42,25 @@ def get_azure_sdk_function(client: Any, function_name: str) -> Callable:
     return func
 
 
+def get_cluster_id_and_nsg_name(resource_group: str,
+                                cluster_name_on_cloud: str) -> Tuple[str, str]:
+    hasher = hashlib.md5(resource_group.encode('utf-8'))
+    unique_id = hasher.hexdigest()[:UNIQUE_ID_LEN]
+    # We use the cluster name + resource group hash as the
+    # unique ID for the cluster, as we need to make sure that
+    # the deployments have unique names during failover.
+    cluster_id = _CLUSTER_ID.format(cluster_name_on_cloud=cluster_name_on_cloud,
+                                    unique_id=unique_id)
+    nsg_name = f'sky-{cluster_id}-nsg'
+    return cluster_id, nsg_name
+
+
 @common.log_function_start_end
 def bootstrap_instances(
         region: str, cluster_name_on_cloud: str,
         config: common.ProvisionConfig) -> common.ProvisionConfig:
     """See sky/provision/__init__.py"""
+    # TODO: use new azure sdk instead of ARM deployment.
     del region  # unused
     provider_config = config.provider_config
     subscription_id = provider_config.get('subscription_id')
@@ -116,12 +131,13 @@ def bootstrap_instances(
 
     logger.info(f'Using cluster name: {cluster_name_on_cloud}')
 
-    hasher = hashlib.md5(provider_config['resource_group'].encode('utf-8'))
-    unique_id = hasher.hexdigest()[:UNIQUE_ID_LEN]
+    cluster_id, nsg_name = get_cluster_id_and_nsg_name(
+        resource_group=provider_config['resource_group'],
+        cluster_name_on_cloud=cluster_name_on_cloud)
     subnet_mask = provider_config.get('subnet_mask')
     if subnet_mask is None:
         # choose a random subnet, skipping most common value of 0
-        random.seed(unique_id)
+        random.seed(cluster_id)
         subnet_mask = f'10.{random.randint(1, 254)}.0.0/16'
     logger.info(f'Using subnet mask: {subnet_mask}')
 
@@ -134,10 +150,10 @@ def bootstrap_instances(
                     'value': subnet_mask
                 },
                 'clusterId': {
-                    # We use the cluster name + resource group hash as the
-                    # unique ID for the cluster, as we need to make sure that
-                    # the deployments have unique names during failover.
-                    'value': f'{cluster_name_on_cloud}-{unique_id}'
+                    'value': cluster_id
+                },
+                'nsgName': {
+                    'value': nsg_name
                 },
             },
         }

sky/provision/azure/instance.py

@@ -2,10 +2,8 @@
 import base64
 import copy
 import enum
-import json
 import logging
 from multiprocessing import pool
-import pathlib
 import time
 import typing
 from typing import Any, Callable, Dict, List, Optional, Tuple
@@ -17,13 +15,16 @@ from sky import status_lib
 from sky.adaptors import azure
 from sky.provision import common
 from sky.provision import constants
+from sky.provision.azure import config as config_lib
 from sky.utils import common_utils
 from sky.utils import subprocess_utils
 from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
     from azure.mgmt import compute as azure_compute
-    from azure.mgmt import resource as azure_resource
+    from azure.mgmt import network as azure_network
+    from azure.mgmt.compute import models as azure_compute_models
+    from azure.mgmt.network import models as azure_network_models
 
 logger = sky_logging.init_logger(__name__)
 
@@ -31,6 +32,8 @@ logger = sky_logging.init_logger(__name__)
 # https://github.com/Azure/azure-sdk-for-python/issues/9422
 azure_logger = logging.getLogger('azure')
 azure_logger.setLevel(logging.WARNING)
+Client = Any
+NetworkSecurityGroup = Any
 
 _RESUME_INSTANCE_TIMEOUT = 480  # 8 minutes
 _RESUME_PER_INSTANCE_TIMEOUT = 120  # 2 minutes
@@ -40,6 +43,10 @@ _WAIT_CREATION_TIMEOUT_SECONDS = 600
 
 _RESOURCE_GROUP_NOT_FOUND_ERROR_MESSAGE = 'ResourceGroupNotFound'
 _POLL_INTERVAL = 1
+# TODO(Doyoung): _LEGACY_NSG_NAME can be remove this after 0.8.0 to ignore
+# legacy nsg names.
+_LEGACY_NSG_NAME = 'ray-{cluster_name_on_cloud}-nsg'
+_SECOND_LEGACY_NSG_NAME = 'sky-{cluster_name_on_cloud}-nsg'
 
 
 class AzureInstanceStatus(enum.Enum):
@@ -184,14 +191,150 @@ def _get_head_instance_id(instances: List) -> Optional[str]:
     return head_instance_id
 
 
-def _create_instances(
-        compute_client: 'azure_compute.ComputeManagementClient',
-        resource_client: 'azure_resource.ResourceManagementClient',
-        cluster_name_on_cloud: str, resource_group: str,
-        provider_config: Dict[str, Any], node_config: Dict[str, Any],
-        tags: Dict[str, str], count: int) -> List:
+def _create_network_interface(
+        network_client: 'azure_network.NetworkManagementClient', vm_name: str,
+        provider_config: Dict[str,
+                              Any]) -> 'azure_network_models.NetworkInterface':
+    network = azure.azure_mgmt_models('network')
+    compute = azure.azure_mgmt_models('compute')
+    logger.info(f'Start creating network interface for {vm_name}...')
+    if provider_config.get('use_internal_ips', False):
+        name = f'{vm_name}-nic-private'
+        ip_config = network.IPConfiguration(
+            name=f'ip-config-private-{vm_name}',
+            subnet=compute.SubResource(id=provider_config['subnet']),
+            private_ip_allocation_method=network.IPAllocationMethod.DYNAMIC)
+    else:
+        name = f'{vm_name}-nic-public'
+        public_ip_address = network.PublicIPAddress(
+            location=provider_config['location'],
+            public_ip_allocation_method='Static',
+            public_ip_address_version='IPv4',
+            sku=network.PublicIPAddressSku(name='Basic', tier='Regional'))
+        ip_poller = network_client.public_ip_addresses.begin_create_or_update(
+            resource_group_name=provider_config['resource_group'],
+            public_ip_address_name=f'{vm_name}-ip',
+            parameters=public_ip_address)
+        logger.info(f'Created public IP address {ip_poller.result().name} '
+                    f'with address {ip_poller.result().ip_address}.')
+        ip_config = network.IPConfiguration(
+            name=f'ip-config-public-{vm_name}',
+            subnet=compute.SubResource(id=provider_config['subnet']),
+            private_ip_allocation_method=network.IPAllocationMethod.DYNAMIC,
+            public_ip_address=network.PublicIPAddress(id=ip_poller.result().id))
+
+    ni_poller = network_client.network_interfaces.begin_create_or_update(
+        resource_group_name=provider_config['resource_group'],
+        network_interface_name=name,
+        parameters=network.NetworkInterface(
+            location=provider_config['location'],
+            ip_configurations=[ip_config],
+            network_security_group=network.NetworkSecurityGroup(
+                id=provider_config['nsg'])))
+    logger.info(f'Created network interface {ni_poller.result().name}.')
+    return ni_poller.result()
+
+
+def _create_vm(
+        compute_client: 'azure_compute.ComputeManagementClient', vm_name: str,
+        node_tags: Dict[str, str], provider_config: Dict[str, Any],
+        node_config: Dict[str, Any],
+        network_interface_id: str) -> 'azure_compute_models.VirtualMachine':
+    compute = azure.azure_mgmt_models('compute')
+    logger.info(f'Start creating VM {vm_name}...')
+    hardware_profile = compute.HardwareProfile(
+        vm_size=node_config['azure_arm_parameters']['vmSize'])
+    network_profile = compute.NetworkProfile(network_interfaces=[
+        compute.NetworkInterfaceReference(id=network_interface_id, primary=True)
+    ])
+    public_key = node_config['azure_arm_parameters']['publicKey']
+    username = node_config['azure_arm_parameters']['adminUsername']
+    os_linux_custom_data = base64.b64encode(
+        node_config['azure_arm_parameters']['cloudInitSetupCommands'].encode(
+            'utf-8')).decode('utf-8')
+    os_profile = compute.OSProfile(
+        admin_username=username,
+        computer_name=vm_name,
+        admin_password=public_key,
+        linux_configuration=compute.LinuxConfiguration(
+            disable_password_authentication=True,
+            ssh=compute.SshConfiguration(public_keys=[
+                compute.SshPublicKey(
+                    path=f'/home/{username}/.ssh/authorized_keys',
+                    key_data=public_key)
+            ])),
+        custom_data=os_linux_custom_data)
+    community_image_id = node_config['azure_arm_parameters'].get(
+        'communityGalleryImageId', None)
+    if community_image_id is not None:
+        # Prioritize using community gallery image if specified.
+        image_reference = compute.ImageReference(
+            community_gallery_image_id=community_image_id)
+        logger.info(
+            f'Used community_image_id: {community_image_id} for VM {vm_name}.')
+    else:
+        image_reference = compute.ImageReference(
+            publisher=node_config['azure_arm_parameters']['imagePublisher'],
+            offer=node_config['azure_arm_parameters']['imageOffer'],
+            sku=node_config['azure_arm_parameters']['imageSku'],
+            version=node_config['azure_arm_parameters']['imageVersion'])
+    storage_profile = compute.StorageProfile(
+        image_reference=image_reference,
+        os_disk=compute.OSDisk(
+            create_option=compute.DiskCreateOptionTypes.FROM_IMAGE,
+            managed_disk=compute.ManagedDiskParameters(
+                storage_account_type=node_config['azure_arm_parameters']
+                ['osDiskTier']),
+            disk_size_gb=node_config['azure_arm_parameters']['osDiskSizeGB']))
+    vm_instance = compute.VirtualMachine(
+        location=provider_config['location'],
+        tags=node_tags,
+        hardware_profile=hardware_profile,
+        os_profile=os_profile,
+        storage_profile=storage_profile,
+        network_profile=network_profile,
+        identity=compute.VirtualMachineIdentity(
+            type='UserAssigned',
+            user_assigned_identities={provider_config['msi']: {}}))
+    vm_poller = compute_client.virtual_machines.begin_create_or_update(
+        resource_group_name=provider_config['resource_group'],
+        vm_name=vm_name,
+        parameters=vm_instance,
+    )
+    # poller.result() will block on async operation until it's done.
+    logger.info(f'Created VM {vm_poller.result().name}.')
+    # Configure driver extension for A10 GPUs. A10 GPUs requires a
+    # special type of drivers which is available at Microsoft HPC
+    # extension. Reference:
+    # https://forums.developer.nvidia.com/t/ubuntu-22-04-installation-driver-error-nvidia-a10/285195/2
+    # This can take more than 20mins for setting up the A10 GPUs
+    if node_config.get('need_nvidia_driver_extension', False):
+        ext_poller = compute_client.virtual_machine_extensions.\
+            begin_create_or_update(
+            resource_group_name=provider_config['resource_group'],
+            vm_name=vm_name,
+            vm_extension_name='NvidiaGpuDriverLinux',
+            extension_parameters=compute.VirtualMachineExtension(
+                location=provider_config['location'],
+                publisher='Microsoft.HpcCompute',
+                type_properties_type='NvidiaGpuDriverLinux',
+                type_handler_version='1.9',
+                auto_upgrade_minor_version=True,
+                settings='{}'))
+        logger.info(
+            f'Created VM extension {ext_poller.result().name} for VM {vm_name}.'
+        )
+    return vm_poller.result()
+
+
+def _create_instances(compute_client: 'azure_compute.ComputeManagementClient',
+                      network_client: 'azure_network.NetworkManagementClient',
+                      cluster_name_on_cloud: str, resource_group: str,
+                      provider_config: Dict[str, Any], node_config: Dict[str,
+                                                                         Any],
+                      tags: Dict[str, str], count: int) -> List:
     vm_id = uuid4().hex[:UNIQUE_ID_LEN]
-    tags = {
+    all_tags = {
         constants.TAG_RAY_CLUSTER_NAME: cluster_name_on_cloud,
         constants.TAG_SKYPILOT_CLUSTER_NAME: cluster_name_on_cloud,
         **constants.WORKER_NODE_TAGS,
@@ -199,83 +342,19 @@ def _create_instances(
         **tags,
     }
     node_tags = node_config['tags'].copy()
-    node_tags.update(tags)
+    node_tags.update(all_tags)
 
-    # load the template file
-    current_path = pathlib.Path(__file__).parent
-    template_path = current_path.joinpath('azure-vm-template.json')
-    with open(template_path, 'r', encoding='utf-8') as template_fp:
-        template = json.load(template_fp)
+    # Create VM instances in parallel.
+    def create_single_instance(vm_i):
+        vm_name = f'{cluster_name_on_cloud}-{vm_id}-{vm_i}'
+        network_interface = _create_network_interface(network_client, vm_name,
+                                                      provider_config)
+        _create_vm(compute_client, vm_name, node_tags, provider_config,
+                   node_config, network_interface.id)
 
-    vm_name = f'{cluster_name_on_cloud}-{vm_id}'
-    use_internal_ips = provider_config.get('use_internal_ips', False)
-
-    template_params = node_config['azure_arm_parameters'].copy()
-    # We don't include 'head' or 'worker' in the VM name as on Azure the VM
-    # name is immutable and we may change the node type for existing VM in the
-    # multi-node cluster, due to manual termination of the head node.
-    template_params['vmName'] = vm_name
-    template_params['provisionPublicIp'] = not use_internal_ips
-    template_params['vmTags'] = node_tags
-    template_params['vmCount'] = count
-    template_params['msi'] = provider_config['msi']
-    template_params['nsg'] = provider_config['nsg']
-    template_params['subnet'] = provider_config['subnet']
-    # In Azure, cloud-init script must be encoded in base64. For more
-    # information, see:
-    # https://learn.microsoft.com/en-us/azure/virtual-machines/custom-data
-    template_params['cloudInitSetupCommands'] = (base64.b64encode(
-        template_params['cloudInitSetupCommands'].encode('utf-8')).decode(
-            'utf-8'))
-
-    if node_config.get('need_nvidia_driver_extension', False):
-        # pylint: disable=line-too-long
-        # Configure driver extension for A10 GPUs. A10 GPUs requires a
-        # special type of drivers which is available at Microsoft HPC
-        # extension. Reference: https://forums.developer.nvidia.com/t/ubuntu-22-04-installation-driver-error-nvidia-a10/285195/2
-        for r in template['resources']:
-            if r['type'] == 'Microsoft.Compute/virtualMachines':
-                # Add a nested extension resource for A10 GPUs
-                r['resources'] = [
-                    {
-                        'type': 'extensions',
-                        'apiVersion': '2015-06-15',
-                        'location': '[variables(\'location\')]',
-                        'dependsOn': [
-                            '[concat(\'Microsoft.Compute/virtualMachines/\', parameters(\'vmName\'), copyIndex())]'
-                        ],
-                        'name': 'NvidiaGpuDriverLinux',
-                        'properties': {
-                            'publisher': 'Microsoft.HpcCompute',
-                            'type': 'NvidiaGpuDriverLinux',
-                            'typeHandlerVersion': '1.9',
-                            'autoUpgradeMinorVersion': True,
-                            'settings': {},
-                        },
-                    },
-                ]
-                break
-
-    parameters = {
-        'properties': {
-            'mode': azure.deployment_mode().incremental,
-            'template': template,
-            'parameters': {
-                key: {
-                    'value': value
-                } for key, value in template_params.items()
-            },
-        }
-    }
-
-    create_or_update = _get_azure_sdk_function(
-        client=resource_client.deployments, function_name='create_or_update')
-    create_or_update(
-        resource_group_name=resource_group,
-        deployment_name=vm_name,
-        parameters=parameters,
-    ).wait()
+    subprocess_utils.run_in_parallel(create_single_instance, range(count))
 
+    # Update disk performance tier
     performance_tier = node_config.get('disk_performance_tier', None)
     if performance_tier is not None:
         disks = compute_client.disks.list_by_resource_group(resource_group)
@@ -286,12 +365,14 @@ def _create_instances(
                 f'az disk update -n {name} -g {resource_group} '
                 f'--set tier={performance_tier}')
 
+    # Validation
     filters = {
         constants.TAG_RAY_CLUSTER_NAME: cluster_name_on_cloud,
         _TAG_SKYPILOT_VM_ID: vm_id
     }
     instances = _filter_instances(compute_client, resource_group, filters)
     assert len(instances) == count, (len(instances), count)
+
     return instances
 
 
@@ -303,7 +384,7 @@ def run_instances(region: str, cluster_name_on_cloud: str,
     resource_group = provider_config['resource_group']
     subscription_id = provider_config['subscription_id']
     compute_client = azure.get_client('compute', subscription_id)
-
+    network_client = azure.get_client('network', subscription_id)
     instances_to_resume = []
     resumed_instance_ids: List[str] = []
     created_instance_ids: List[str] = []
@@ -439,12 +520,11 @@ def run_instances(region: str, cluster_name_on_cloud: str,
     to_start_count -= len(resumed_instance_ids)
 
     if to_start_count > 0:
-        resource_client = azure.get_client('resource', subscription_id)
         logger.debug(f'run_instances: Creating {to_start_count} instances.')
         try:
             created_instances = _create_instances(
                 compute_client=compute_client,
-                resource_client=resource_client,
+                network_client=network_client,
                 cluster_name_on_cloud=cluster_name_on_cloud,
                 resource_group=resource_group,
                 provider_config=provider_config,
@@ -722,6 +802,32 @@ def query_instances(
     return statuses
 
 
+# TODO(Doyoung): _get_cluster_nsg can be remove this after 0.8.0 to ignore
+# legacy nsg names.
+def _get_cluster_nsg(network_client: Client, resource_group: str,
+                     cluster_name_on_cloud: str) -> NetworkSecurityGroup:
+    """Retrieve the NSG associated with the given name of the cluster."""
+    list_network_security_groups = _get_azure_sdk_function(
+        client=network_client.network_security_groups, function_name='list')
+    legacy_nsg_name = _LEGACY_NSG_NAME.format(
+        cluster_name_on_cloud=cluster_name_on_cloud)
+    second_legacy_nsg_name = _SECOND_LEGACY_NSG_NAME.format(
+        cluster_name_on_cloud=cluster_name_on_cloud)
+    _, nsg_name = config_lib.get_cluster_id_and_nsg_name(
+        resource_group=resource_group,
+        cluster_name_on_cloud=cluster_name_on_cloud)
+    possible_nsg_names = [nsg_name, legacy_nsg_name, second_legacy_nsg_name]
+    for nsg in list_network_security_groups(resource_group):
+        if nsg.name in possible_nsg_names:
+            return nsg
+
+    # Raise an error if no matching NSG is found
+    raise ValueError('Failed to find a matching NSG for cluster '
+                     f'{cluster_name_on_cloud!r} in resource group '
+                     f'{resource_group!r}. Expected NSG names were: '
+                     f'{possible_nsg_names}.')
+
+
 def open_ports(
     cluster_name_on_cloud: str,
     ports: List[str],
@@ -736,58 +842,66 @@ def open_ports(
     update_network_security_groups = _get_azure_sdk_function(
         client=network_client.network_security_groups,
         function_name='create_or_update')
-    list_network_security_groups = _get_azure_sdk_function(
-        client=network_client.network_security_groups, function_name='list')
-    for nsg in list_network_security_groups(resource_group):
-        try:
-            # Wait the NSG creation to be finished before opening a port. The
-            # cluster provisioning triggers the NSG creation, but it may not be
-            # finished yet.
-            backoff = common_utils.Backoff(max_backoff_factor=1)
-            start_time = time.time()
-            while True:
-                if nsg.provisioning_state not in ['Creating', 'Updating']:
-                    break
-                if time.time() - start_time > _WAIT_CREATION_TIMEOUT_SECONDS:
-                    logger.warning(
-                        f'Fails to wait for the creation of NSG {nsg.name} in '
-                        f'{resource_group} within '
-                        f'{_WAIT_CREATION_TIMEOUT_SECONDS} seconds. '
-                        'Skip this NSG.')
-                backoff_time = backoff.current_backoff()
-                logger.info(f'NSG {nsg.name} is not created yet. Waiting for '
-                            f'{backoff_time} seconds before checking again.')
-                time.sleep(backoff_time)
-
-            # Azure NSG rules have a priority field that determines the order
-            # in which they are applied. The priority must be unique across
-            # all inbound rules in one NSG.
-            priority = max(rule.priority
-                           for rule in nsg.security_rules
-                           if rule.direction == 'Inbound') + 1
-            nsg.security_rules.append(
-                azure.create_security_rule(
-                    name=f'sky-ports-{cluster_name_on_cloud}-{priority}',
-                    priority=priority,
-                    protocol='Tcp',
-                    access='Allow',
-                    direction='Inbound',
-                    source_address_prefix='*',
-                    source_port_range='*',
-                    destination_address_prefix='*',
-                    destination_port_ranges=ports,
-                ))
-            poller = update_network_security_groups(resource_group, nsg.name,
-                                                    nsg)
-            poller.wait()
-            if poller.status() != 'Succeeded':
+
+    try:
+        # Wait for the NSG creation to be finished before opening a port. The
+        # cluster provisioning triggers the NSG creation, but it may not be
+        # finished yet.
+        backoff = common_utils.Backoff(max_backoff_factor=1)
+        start_time = time.time()
+        while True:
+            nsg = _get_cluster_nsg(network_client, resource_group,
+                                   cluster_name_on_cloud)
+            if nsg.provisioning_state not in ['Creating', 'Updating']:
+                break
+            if time.time() - start_time > _WAIT_CREATION_TIMEOUT_SECONDS:
                 with ux_utils.print_exception_no_traceback():
-                    raise ValueError(f'Failed to open ports {ports} in NSG '
-                                     f'{nsg.name}: {poller.status()}')
-        except azure.exceptions().HttpResponseError as e:
+                    raise TimeoutError(
+                        f'Timed out while waiting for the Network '
+                        f'Security Group {nsg.name!r} to be ready for '
+                        f'cluster {cluster_name_on_cloud!r} in '
+                        f'resource group {resource_group!r}. The NSG '
+                        f'did not reach a stable state '
+                        '(Creating/Updating) within the allocated '
+                        f'{_WAIT_CREATION_TIMEOUT_SECONDS} seconds. '
+                        'Consequently, the operation to open ports '
+                        f'{ports} failed.')
+
+            backoff_time = backoff.current_backoff()
+            logger.info(f'NSG {nsg.name} is not created yet. Waiting for '
+                        f'{backoff_time} seconds before checking again.')
+            time.sleep(backoff_time)
+
+        # Azure NSG rules have a priority field that determines the order
+        # in which they are applied. The priority must be unique across
+        # all inbound rules in one NSG.
+        priority = max(rule.priority
+                       for rule in nsg.security_rules
+                       if rule.direction == 'Inbound') + 1
+        nsg.security_rules.append(
+            azure.create_security_rule(
+                name=f'sky-ports-{cluster_name_on_cloud}-{priority}',
+                priority=priority,
+                protocol='Tcp',
+                access='Allow',
+                direction='Inbound',
+                source_address_prefix='*',
+                source_port_range='*',
+                destination_address_prefix='*',
+                destination_port_ranges=ports,
+            ))
+        poller = update_network_security_groups(resource_group, nsg.name, nsg)
+        poller.wait()
+        if poller.status() != 'Succeeded':
             with ux_utils.print_exception_no_traceback():
-                raise ValueError(
-                    f'Failed to open ports {ports} in NSG {nsg.name}.') from e
+                raise ValueError(f'Failed to open ports {ports} in NSG '
+                                 f'{nsg.name}: {poller.status()}')
+
+    except azure.exceptions().HttpResponseError as e:
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(f'Failed to open ports {ports} in NSG for cluster '
+                             f'{cluster_name_on_cloud!r} within resource group '
+                             f'{resource_group!r}.') from e
 
 
 def cleanup_ports(

sky/provision/kubernetes/instance.py

@@ -18,6 +18,7 @@ from sky.provision.kubernetes import utils as kubernetes_utils
 from sky.utils import command_runner
 from sky.utils import common_utils
 from sky.utils import kubernetes_enums
+from sky.utils import subprocess_utils
 from sky.utils import ux_utils
 
 POLL_INTERVAL = 2
@@ -398,8 +399,7 @@ def _setup_ssh_in_pods(namespace: str, context: Optional[str],
         # See https://www.educative.io/answers/error-mesg-ttyname-failed-inappropriate-ioctl-for-device  # pylint: disable=line-too-long
         '$(prefix_cmd) sed -i "s/mesg n/tty -s \\&\\& mesg n/" ~/.profile;')
 
-    # TODO(romilb): Parallelize the setup of SSH in pods for multi-node clusters
-    for new_node in new_nodes:
+    def _setup_ssh_thread(new_node):
         pod_name = new_node.metadata.name
         runner = command_runner.KubernetesCommandRunner(
             ((namespace, context), pod_name))
@@ -411,6 +411,8 @@ def _setup_ssh_in_pods(namespace: str, context: Optional[str],
                                      stdout)
         logger.info(f'{"-"*20}End: Set up SSH in pod {pod_name!r} {"-"*20}')
 
+    subprocess_utils.run_in_parallel(_setup_ssh_thread, new_nodes)
+
 
 def _label_pod(namespace: str, context: Optional[str], pod_name: str,
                label: Dict[str, str]) -> None: