skypilot-nightly 1.0.0.dev2024053101__py3-none-any.whl → 1.0.0.dev2025022801__py3-none-any.whl

sky/provision/azure/instance.py

@@ -1,24 +1,142 @@
 """Azure instance provisioning."""
+import base64
+import copy
+import enum
 import logging
-from typing import Any, Callable, Dict, List, Optional
+from multiprocessing import pool
+import time
+import typing
+from typing import Any, Callable, Dict, List, Optional, Tuple
+from uuid import uuid4
 
+from sky import exceptions
 from sky import sky_logging
 from sky.adaptors import azure
+from sky.provision import common
+from sky.provision import constants
+from sky.utils import common_utils
+from sky.utils import status_lib
+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 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__)
 
 # Suppress noisy logs from Azure SDK. Reference:
 # 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
+UNIQUE_ID_LEN = 4
+_TAG_SKYPILOT_VM_ID = 'skypilot-vm-id'
+_WAIT_CREATION_TIMEOUT_SECONDS = 600
+
+_RESOURCE_MANAGED_IDENTITY_TYPE = (
+    'Microsoft.ManagedIdentity/userAssignedIdentities')
+_RESOURCE_NETWORK_SECURITY_GROUP_TYPE = (
+    'Microsoft.Network/networkSecurityGroups')
+_RESOURCE_VIRTUAL_NETWORK_TYPE = 'Microsoft.Network/virtualNetworks'
+_RESOURCE_PUBLIC_IP_ADDRESS_TYPE = 'Microsoft.Network/publicIPAddresses'
+_RESOURCE_VIRTUAL_MACHINE_TYPE = 'Microsoft.Compute/virtualMachines'
+_RESOURCE_NETWORK_INTERFACE_TYPE = 'Microsoft.Network/networkInterfaces'
+
+_RESOURCE_GROUP_NOT_FOUND_ERROR_MESSAGE = 'ResourceGroupNotFound'
+_POLL_INTERVAL = 1
+
 
-# Tag uniquely identifying all nodes of a cluster
-TAG_RAY_CLUSTER_NAME = 'ray-cluster-name'
-TAG_RAY_NODE_KIND = 'ray-node-type'
+class AzureInstanceStatus(enum.Enum):
+    """Statuses enum for Azure instances with power and provisioning states."""
+    PENDING = 'pending'
+    RUNNING = 'running'
+    STOPPING = 'stopping'
+    STOPPED = 'stopped'
+    DELETING = 'deleting'
 
+    @classmethod
+    def power_state_map(cls) -> Dict[str, 'AzureInstanceStatus']:
+        return {
+            'starting': cls.PENDING,
+            'running': cls.RUNNING,
+            # 'stopped' in Azure means Stopped (Allocated), which still bills
+            # for the VM.
+            'stopping': cls.STOPPING,
+            'stopped': cls.STOPPED,
+            # 'VM deallocated' in Azure means Stopped (Deallocated), which does
+            # not bill for the VM.
+            'deallocating': cls.STOPPING,
+            'deallocated': cls.STOPPED,
+        }
 
-def get_azure_sdk_function(client: Any, function_name: str) -> Callable:
+    @classmethod
+    def provisioning_state_map(cls) -> Dict[str, 'AzureInstanceStatus']:
+        return {
+            'Creating': cls.PENDING,
+            'Updating': cls.PENDING,
+            'Failed': cls.PENDING,
+            'Migrating': cls.PENDING,
+            'Deleting': cls.DELETING,
+            # Succeeded in provisioning state means the VM is provisioned but
+            # not necessarily running. The caller should further check the
+            # power state to determine the actual VM status.
+            'Succeeded': cls.RUNNING,
+        }
+
+    @classmethod
+    def cluster_status_map(
+            cls
+    ) -> Dict['AzureInstanceStatus', Optional[status_lib.ClusterStatus]]:
+        return {
+            cls.PENDING: status_lib.ClusterStatus.INIT,
+            cls.RUNNING: status_lib.ClusterStatus.UP,
+            cls.STOPPING: status_lib.ClusterStatus.STOPPED,
+            cls.STOPPED: status_lib.ClusterStatus.STOPPED,
+            cls.DELETING: None,
+        }
+
+    @classmethod
+    def from_raw_states(cls, provisioning_state: str,
+                        power_state: Optional[str]) -> 'AzureInstanceStatus':
+        provisioning_state_map = cls.provisioning_state_map()
+        power_state_map = cls.power_state_map()
+        status = None
+        if power_state is None:
+            if provisioning_state not in provisioning_state_map:
+                with ux_utils.print_exception_no_traceback():
+                    raise exceptions.ClusterStatusFetchingError(
+                        'Failed to parse status from Azure response: '
+                        f'{provisioning_state}')
+            status = provisioning_state_map[provisioning_state]
+        if status is None or status == cls.RUNNING:
+            # We should further check the power state to determine the actual
+            # VM status.
+            if power_state not in power_state_map:
+                with ux_utils.print_exception_no_traceback():
+                    raise exceptions.ClusterStatusFetchingError(
+                        'Failed to parse status from Azure response: '
+                        f'{power_state}.')
+            status = power_state_map[power_state]
+        if status is None:
+            with ux_utils.print_exception_no_traceback():
+                raise exceptions.ClusterStatusFetchingError(
+                    'Failed to parse status from Azure response: '
+                    f'provisioning state ({provisioning_state}), '
+                    f'power state ({power_state})')
+        return status
+
+    def to_cluster_status(self) -> Optional[status_lib.ClusterStatus]:
+        return self.cluster_status_map().get(self)
+
+
+def _get_azure_sdk_function(client: Any, function_name: str) -> Callable:
     """Retrieve a callable function from Azure SDK client object.
 
     Newer versions of the various client SDKs renamed function names to
@@ -35,6 +153,835 @@ def get_azure_sdk_function(client: Any, function_name: str) -> Callable:
     return func
 
 
+def _get_instance_ips(network_client, vm, resource_group: str,
+                      use_internal_ips: bool) -> Tuple[str, Optional[str]]:
+    nic_id = vm.network_profile.network_interfaces[0].id
+    nic_name = nic_id.split('/')[-1]
+    nic = network_client.network_interfaces.get(
+        resource_group_name=resource_group,
+        network_interface_name=nic_name,
+    )
+    ip_config = nic.ip_configurations[0]
+
+    external_ip = None
+    if not use_internal_ips:
+        public_ip_id = ip_config.public_ip_address.id
+        public_ip_name = public_ip_id.split('/')[-1]
+        public_ip = network_client.public_ip_addresses.get(
+            resource_group_name=resource_group,
+            public_ip_address_name=public_ip_name,
+        )
+        external_ip = public_ip.ip_address
+
+    internal_ip = ip_config.private_ip_address
+
+    return (internal_ip, external_ip)
+
+
+def _get_head_instance_id(instances: List) -> Optional[str]:
+    head_instance_id = None
+    head_node_tags = tuple(constants.HEAD_NODE_TAGS.items())
+    for inst in instances:
+        for k, v in inst.tags.items():
+            if (k, v) in head_node_tags:
+                if head_instance_id is not None:
+                    logger.warning(
+                        'There are multiple head nodes in the cluster '
+                        f'(current head instance id: {head_instance_id}, '
+                        f'newly discovered id: {inst.name}). It is likely '
+                        f'that something goes wrong.')
+                head_instance_id = inst.name
+                break
+    return head_instance_id
+
+
+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,
+            delete_option=compute.DiskDeleteOptionTypes.DELETE,
+            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']: {}}),
+        priority=node_config['azure_arm_parameters'].get('priority', None))
+    vm_poller = compute_client.virtual_machines.begin_create_or_update(
+        resource_group_name=provider_config['resource_group'],
+        vm_name=vm_name,
+        parameters=vm_instance,
+    )
+    # This line will block until the VM is created or the operation times out.
+    vm = vm_poller.result()
+    logger.info(f'Created VM {vm.name}.')
+    return vm
+
+
+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]
+    all_tags = {
+        constants.TAG_RAY_CLUSTER_NAME: cluster_name_on_cloud,
+        constants.TAG_SKYPILOT_CLUSTER_NAME: cluster_name_on_cloud,
+        **constants.WORKER_NODE_TAGS,
+        _TAG_SKYPILOT_VM_ID: vm_id,
+        **tags,
+    }
+    node_tags = node_config['tags'].copy()
+    node_tags.update(all_tags)
+
+    # 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)
+
+    subprocess_utils.run_in_parallel(create_single_instance, list(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)
+        for disk in disks:
+            name = disk.name
+            # TODO(tian): Investigate if we can use Python SDK to update this.
+            subprocess_utils.run_no_outputs(
+                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
+
+
+def run_instances(region: str, cluster_name_on_cloud: str,
+                  config: common.ProvisionConfig) -> common.ProvisionRecord:
+    """See sky/provision/__init__.py"""
+    # TODO(zhwu): This function is too long. We should refactor it.
+    provider_config = config.provider_config
+    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] = []
+
+    # sort tags by key to support deterministic unit test stubbing
+    tags = dict(sorted(copy.deepcopy(config.tags).items()))
+    filters = {constants.TAG_RAY_CLUSTER_NAME: cluster_name_on_cloud}
+
+    non_deleting_states = (set(AzureInstanceStatus) -
+                           {AzureInstanceStatus.DELETING})
+    existing_instances = _filter_instances(
+        compute_client,
+        tag_filters=filters,
+        resource_group=resource_group,
+        status_filters=list(non_deleting_states),
+    )
+    logger.debug(
+        f'run_instances: Found {[inst.name for inst in existing_instances]} '
+        'existing instances in cluster.')
+    existing_instances.sort(key=lambda x: x.name)
+
+    pending_instances = []
+    running_instances = []
+    stopping_instances = []
+    stopped_instances = []
+
+    for instance in existing_instances:
+        status = _get_instance_status(compute_client, instance, resource_group)
+        logger.debug(
+            f'run_instances: Instance {instance.name} has status {status}.')
+
+        if status == AzureInstanceStatus.RUNNING:
+            running_instances.append(instance)
+        elif status == AzureInstanceStatus.STOPPED:
+            stopped_instances.append(instance)
+        elif status == AzureInstanceStatus.STOPPING:
+            stopping_instances.append(instance)
+        elif status == AzureInstanceStatus.PENDING:
+            pending_instances.append(instance)
+
+    def _create_instance_tag(target_instance, is_head: bool = True) -> str:
+        new_instance_tags = (constants.HEAD_NODE_TAGS
+                             if is_head else constants.WORKER_NODE_TAGS)
+
+        tags = target_instance.tags
+        tags.update(new_instance_tags)
+
+        update = _get_azure_sdk_function(compute_client.virtual_machines,
+                                         'update')
+        update(resource_group, target_instance.name, parameters={'tags': tags})
+        return target_instance.name
+
+    head_instance_id = _get_head_instance_id(existing_instances)
+    if head_instance_id is None:
+        if running_instances:
+            head_instance_id = _create_instance_tag(running_instances[0])
+        elif pending_instances:
+            head_instance_id = _create_instance_tag(pending_instances[0])
+
+    if config.resume_stopped_nodes and len(existing_instances) > config.count:
+        raise RuntimeError(
+            'The number of pending/running/stopped/stopping '
+            f'instances combined ({len(existing_instances)}) in '
+            f'cluster "{cluster_name_on_cloud}" is greater than the '
+            f'number requested by the user ({config.count}). '
+            'This is likely a resource leak. '
+            'Use "sky down" to terminate the cluster.')
+
+    to_start_count = config.count - len(pending_instances) - len(
+        running_instances)
+
+    if to_start_count < 0:
+        raise RuntimeError(
+            'The number of running+pending instances '
+            f'({config.count - to_start_count}) in cluster '
+            f'"{cluster_name_on_cloud}" is greater than the number '
+            f'requested by the user ({config.count}). '
+            'This is likely a resource leak. '
+            'Use "sky down" to terminate the cluster.')
+
+    if config.resume_stopped_nodes and to_start_count > 0 and (
+            stopping_instances or stopped_instances):
+        time_start = time.time()
+        if stopping_instances:
+            plural = 's' if len(stopping_instances) > 1 else ''
+            verb = 'are' if len(stopping_instances) > 1 else 'is'
+            # TODO(zhwu): double check the correctness of the following on Azure
+            logger.warning(
+                f'Instance{plural} {[inst.name for inst in stopping_instances]}'
+                f' {verb} still in STOPPING state on Azure. It can only be '
+                'resumed after it is fully STOPPED. Waiting ...')
+        while (stopping_instances and
+               to_start_count > len(stopped_instances) and
+               time.time() - time_start < _RESUME_INSTANCE_TIMEOUT):
+            inst = stopping_instances.pop(0)
+            per_instance_time_start = time.time()
+            while (time.time() - per_instance_time_start <
+                   _RESUME_PER_INSTANCE_TIMEOUT):
+                status = _get_instance_status(compute_client, inst,
+                                              resource_group)
+                if status == AzureInstanceStatus.STOPPED:
+                    break
+                time.sleep(1)
+            else:
+                logger.warning(
+                    f'Instance {inst.name} is still in stopping state '
+                    f'(Timeout: {_RESUME_PER_INSTANCE_TIMEOUT}). '
+                    'Retrying ...')
+                stopping_instances.append(inst)
+                time.sleep(5)
+                continue
+            stopped_instances.append(inst)
+        if stopping_instances and to_start_count > len(stopped_instances):
+            msg = ('Timeout for waiting for existing instances '
+                   f'{stopping_instances} in STOPPING state to '
+                   'be STOPPED before restarting them. Please try again later.')
+            logger.error(msg)
+            raise RuntimeError(msg)
+
+        instances_to_resume = stopped_instances[:to_start_count]
+        instances_to_resume.sort(key=lambda x: x.name)
+        instances_to_resume_ids = [t.name for t in instances_to_resume]
+        logger.debug('run_instances: Resuming stopped instances '
+                     f'{instances_to_resume_ids}.')
+        start_virtual_machine = _get_azure_sdk_function(
+            compute_client.virtual_machines, 'start')
+        with pool.ThreadPool() as p:
+            p.starmap(
+                start_virtual_machine,
+                [(resource_group, inst.name) for inst in instances_to_resume])
+        resumed_instance_ids = instances_to_resume_ids
+
+    to_start_count -= len(resumed_instance_ids)
+
+    if to_start_count > 0:
+        logger.debug(f'run_instances: Creating {to_start_count} instances.')
+        try:
+            created_instances = _create_instances(
+                compute_client=compute_client,
+                network_client=network_client,
+                cluster_name_on_cloud=cluster_name_on_cloud,
+                resource_group=resource_group,
+                provider_config=provider_config,
+                node_config=config.node_config,
+                tags=tags,
+                count=to_start_count)
+        except Exception as e:
+            err_message = common_utils.format_exception(
+                e, use_bracket=True).replace('\n', ' ')
+            logger.error(f'Failed to create instances: {err_message}')
+            raise
+        created_instance_ids = [inst.name for inst in created_instances]
+
+    non_running_instance_statuses = list(
+        set(AzureInstanceStatus) - {AzureInstanceStatus.RUNNING})
+    start = time.time()
+    while True:
+        # Wait for all instances to be in running state
+        instances = _filter_instances(
+            compute_client,
+            resource_group,
+            filters,
+            status_filters=non_running_instance_statuses,
+            included_instances=created_instance_ids + resumed_instance_ids)
+        if not instances:
+            break
+        if time.time() - start > _WAIT_CREATION_TIMEOUT_SECONDS:
+            raise TimeoutError(
+                'run_instances: Timed out waiting for Azure instances to be '
+                f'running: {instances}')
+        logger.debug(f'run_instances: Waiting for {len(instances)} instances '
+                     'in PENDING status.')
+        time.sleep(_POLL_INTERVAL)
+
+    running_instances = _filter_instances(
+        compute_client,
+        resource_group,
+        filters,
+        status_filters=[AzureInstanceStatus.RUNNING])
+    head_instance_id = _get_head_instance_id(running_instances)
+    instances_to_tag = copy.copy(running_instances)
+    if head_instance_id is None:
+        head_instance_id = _create_instance_tag(instances_to_tag[0])
+        instances_to_tag = instances_to_tag[1:]
+    else:
+        instances_to_tag = [
+            inst for inst in instances_to_tag if inst.name != head_instance_id
+        ]
+
+    if instances_to_tag:
+        # Tag the instances in case the old resumed instances are not correctly
+        # tagged.
+        with pool.ThreadPool() as p:
+            p.starmap(
+                _create_instance_tag,
+                # is_head=False for all wokers.
+                [(inst, False) for inst in instances_to_tag])
+
+    assert head_instance_id is not None, head_instance_id
+    return common.ProvisionRecord(
+        provider_name='azure',
+        region=region,
+        zone=None,
+        cluster_name=cluster_name_on_cloud,
+        head_instance_id=head_instance_id,
+        created_instance_ids=created_instance_ids,
+        resumed_instance_ids=resumed_instance_ids,
+    )
+
+
+def wait_instances(region: str, cluster_name_on_cloud: str,
+                   state: Optional[status_lib.ClusterStatus]) -> None:
+    """See sky/provision/__init__.py"""
+    del region, cluster_name_on_cloud, state
+    # We already wait for the instances to be running in run_instances.
+    # So we don't need to wait here.
+
+
+def get_cluster_info(
+        region: str,
+        cluster_name_on_cloud: str,
+        provider_config: Optional[Dict[str, Any]] = None) -> common.ClusterInfo:
+    """See sky/provision/__init__.py"""
+    del region
+    filters = {constants.TAG_RAY_CLUSTER_NAME: cluster_name_on_cloud}
+    assert provider_config is not None, (cluster_name_on_cloud, provider_config)
+    resource_group = provider_config['resource_group']
+    subscription_id = provider_config.get('subscription_id',
+                                          azure.get_subscription_id())
+    compute_client = azure.get_client('compute', subscription_id)
+    network_client = azure.get_client('network', subscription_id)
+
+    running_instances = _filter_instances(
+        compute_client,
+        resource_group,
+        filters,
+        status_filters=[AzureInstanceStatus.RUNNING])
+    head_instance_id = _get_head_instance_id(running_instances)
+
+    instances = {}
+    use_internal_ips = provider_config.get('use_internal_ips', False)
+    for inst in running_instances:
+        internal_ip, external_ip = _get_instance_ips(network_client, inst,
+                                                     resource_group,
+                                                     use_internal_ips)
+        instances[inst.name] = [
+            common.InstanceInfo(
+                instance_id=inst.name,
+                internal_ip=internal_ip,
+                external_ip=external_ip,
+                tags=inst.tags,
+            )
+        ]
+    instances = dict(sorted(instances.items(), key=lambda x: x[0]))
+    return common.ClusterInfo(
+        provider_name='azure',
+        head_instance_id=head_instance_id,
+        instances=instances,
+        provider_config=provider_config,
+    )
+
+
+def stop_instances(
+    cluster_name_on_cloud: str,
+    provider_config: Optional[Dict[str, Any]] = None,
+    worker_only: bool = False,
+) -> None:
+    """See sky/provision/__init__.py"""
+    assert provider_config is not None, (cluster_name_on_cloud, provider_config)
+
+    subscription_id = provider_config['subscription_id']
+    resource_group = provider_config['resource_group']
+    compute_client = azure.get_client('compute', subscription_id)
+    tag_filters = {constants.TAG_RAY_CLUSTER_NAME: cluster_name_on_cloud}
+    if worker_only:
+        tag_filters[constants.TAG_RAY_NODE_KIND] = 'worker'
+
+    nodes = _filter_instances(compute_client, resource_group, tag_filters)
+    stop_virtual_machine = _get_azure_sdk_function(
+        client=compute_client.virtual_machines, function_name='deallocate')
+    with pool.ThreadPool() as p:
+        p.starmap(stop_virtual_machine,
+                  [(resource_group, node.name) for node in nodes])
+
+
+def terminate_instances(
+    cluster_name_on_cloud: str,
+    provider_config: Optional[Dict[str, Any]] = None,
+    worker_only: bool = False,
+) -> None:
+    """See sky/provision/__init__.py"""
+    assert provider_config is not None, (cluster_name_on_cloud, provider_config)
+    # TODO(zhwu): check the following. Also, seems we can directly force
+    # delete a resource group.
+    subscription_id = provider_config['subscription_id']
+    resource_group = provider_config['resource_group']
+    if worker_only:
+        compute_client = azure.get_client('compute', subscription_id)
+        delete_virtual_machine = _get_azure_sdk_function(
+            client=compute_client.virtual_machines, function_name='delete')
+        filters = {
+            constants.TAG_RAY_CLUSTER_NAME: cluster_name_on_cloud,
+            constants.TAG_RAY_NODE_KIND: 'worker'
+        }
+        nodes = _filter_instances(compute_client, resource_group, filters)
+        with pool.ThreadPool() as p:
+            p.starmap(delete_virtual_machine,
+                      [(resource_group, node.name) for node in nodes])
+        return
+
+    assert provider_config is not None, cluster_name_on_cloud
+
+    use_external_resource_group = provider_config.get(
+        'use_external_resource_group', False)
+    # When user specified resource group through config.yaml to create a VM, we
+    # cannot remove the entire resource group as it may contain other resources
+    # unrelated to this VM being removed.
+    if use_external_resource_group:
+        delete_vm_and_attached_resources(subscription_id, resource_group,
+                                         cluster_name_on_cloud)
+    else:
+        # For SkyPilot default resource groups, delete entire resource group.
+        # This automatically terminates all resources within, including VMs
+        resource_group_client = azure.get_client('resource', subscription_id)
+        delete_resource_group = _get_azure_sdk_function(
+            client=resource_group_client.resource_groups,
+            function_name='delete')
+        try:
+            delete_resource_group(resource_group, force_deletion_types=None)
+        except azure.exceptions().ResourceNotFoundError as e:
+            if 'ResourceGroupNotFound' in str(e):
+                logger.warning(
+                    f'Resource group {resource_group} not found. Skip '
+                    'terminating it.')
+                return
+            raise
+
+
+def _get_instance_status(
+        compute_client: 'azure_compute.ComputeManagementClient', vm,
+        resource_group: str) -> Optional[AzureInstanceStatus]:
+    try:
+        instance = compute_client.virtual_machines.instance_view(
+            resource_group_name=resource_group, vm_name=vm.name)
+    except azure.exceptions().ResourceNotFoundError as e:
+        if 'ResourceNotFound' in str(e):
+            return None
+        raise
+    provisioning_state = vm.provisioning_state
+    instance_dict = instance.as_dict()
+    for status in instance_dict['statuses']:
+        code_state = status['code'].split('/')
+        # It is possible that sometimes the 'code' is empty string, and we
+        # should skip them.
+        if len(code_state) != 2:
+            continue
+        code, state = code_state
+        # skip provisioning status
+        if code == 'PowerState':
+            return AzureInstanceStatus.from_raw_states(provisioning_state,
+                                                       state)
+    return AzureInstanceStatus.from_raw_states(provisioning_state, None)
+
+
+def _filter_instances(
+    compute_client: 'azure_compute.ComputeManagementClient',
+    resource_group: str,
+    tag_filters: Dict[str, str],
+    status_filters: Optional[List[AzureInstanceStatus]] = None,
+    included_instances: Optional[List[str]] = None,
+) -> List['azure_compute.models.VirtualMachine']:
+
+    def match_tags(vm):
+        for k, v in tag_filters.items():
+            if vm.tags.get(k) != v:
+                return False
+        return True
+
+    try:
+        list_virtual_machines = _get_azure_sdk_function(
+            client=compute_client.virtual_machines, function_name='list')
+        vms = list_virtual_machines(resource_group_name=resource_group)
+        nodes = list(filter(match_tags, vms))
+    except azure.exceptions().ResourceNotFoundError as e:
+        if _RESOURCE_GROUP_NOT_FOUND_ERROR_MESSAGE in str(e):
+            return []
+        raise
+    if status_filters is not None:
+        nodes = [
+            node for node in nodes if _get_instance_status(
+                compute_client, node, resource_group) in status_filters
+        ]
+    if included_instances:
+        nodes = [node for node in nodes if node.name in included_instances]
+    return nodes
+
+
+def _delete_nic_with_retries(network_client,
+                             resource_group,
+                             nic_name,
+                             max_retries=15,
+                             retry_interval=20):
+    """Delete a NIC with retries.
+
+    When a VM is created, its NIC is reserved for 180 seconds, preventing its
+    immediate deletion. If the NIC is in this reserved state, we must retry
+    deletion with intervals until the reservation expires. This situation
+    commonly arises if a VM termination is followed by a failover to another
+    region due to provisioning failures.
+    """
+    delete_network_interfaces = _get_azure_sdk_function(
+        client=network_client.network_interfaces, function_name='begin_delete')
+    for _ in range(max_retries):
+        try:
+            delete_network_interfaces(resource_group_name=resource_group,
+                                      network_interface_name=nic_name).result()
+            return
+        except azure.exceptions().HttpResponseError as e:
+            if 'NicReservedForAnotherVm' in str(e):
+                # Retry when deletion fails with reserved NIC.
+                logger.warning(f'NIC {nic_name} is reserved. '
+                               f'Retrying in {retry_interval} seconds...')
+                time.sleep(retry_interval)
+            else:
+                raise e
+    logger.error(
+        f'Failed to delete NIC {nic_name} after {max_retries} attempts.')
+
+
+def delete_vm_and_attached_resources(subscription_id: str, resource_group: str,
+                                     cluster_name_on_cloud: str) -> None:
+    """Removes VM with attached resources and Deployments.
+
+    This function deletes a virtual machine and its associated resources
+    (public IP addresses, virtual networks, managed identities, network
+    interface and network security groups) that match cluster_name_on_cloud.
+    There is one attached resources that is not removed within this
+    method: OS disk. It is configured to be deleted when VM is terminated while
+    setting up storage profile from _create_vm.
+
+    Args:
+        subscription_id: The Azure subscription ID.
+        resource_group: The name of the resource group.
+        cluster_name_on_cloud: The name of the cluster to filter resources.
+    """
+    resource_client = azure.get_client('resource', subscription_id)
+    try:
+        list_resources = _get_azure_sdk_function(
+            client=resource_client.resources,
+            function_name='list_by_resource_group')
+        resources = list(list_resources(resource_group))
+    except azure.exceptions().ResourceNotFoundError as e:
+        if _RESOURCE_GROUP_NOT_FOUND_ERROR_MESSAGE in str(e):
+            return
+        raise
+
+    filtered_resources: Dict[str, List[str]] = {
+        _RESOURCE_VIRTUAL_MACHINE_TYPE: [],
+        _RESOURCE_MANAGED_IDENTITY_TYPE: [],
+        _RESOURCE_NETWORK_SECURITY_GROUP_TYPE: [],
+        _RESOURCE_VIRTUAL_NETWORK_TYPE: [],
+        _RESOURCE_PUBLIC_IP_ADDRESS_TYPE: [],
+        _RESOURCE_NETWORK_INTERFACE_TYPE: []
+    }
+
+    for resource in resources:
+        if (resource.type in filtered_resources and
+                cluster_name_on_cloud in resource.name):
+            filtered_resources[resource.type].append(resource.name)
+
+    network_client = azure.get_client('network', subscription_id)
+    msi_client = azure.get_client('msi', subscription_id)
+    compute_client = azure.get_client('compute', subscription_id)
+    auth_client = azure.get_client('authorization', subscription_id)
+
+    delete_virtual_machine = _get_azure_sdk_function(
+        client=compute_client.virtual_machines, function_name='delete')
+    delete_public_ip_addresses = _get_azure_sdk_function(
+        client=network_client.public_ip_addresses, function_name='begin_delete')
+    delete_virtual_networks = _get_azure_sdk_function(
+        client=network_client.virtual_networks, function_name='begin_delete')
+    delete_managed_identity = _get_azure_sdk_function(
+        client=msi_client.user_assigned_identities, function_name='delete')
+    delete_network_security_group = _get_azure_sdk_function(
+        client=network_client.network_security_groups,
+        function_name='begin_delete')
+    delete_role_assignment = _get_azure_sdk_function(
+        client=auth_client.role_assignments, function_name='delete')
+
+    for vm_name in filtered_resources[_RESOURCE_VIRTUAL_MACHINE_TYPE]:
+        try:
+            # Before removing Network Interface, we need to wait for the VM to
+            # be completely removed with .result() so the dependency of VM on
+            # Network Interface is disassociated. This takes abour ~30s.
+            delete_virtual_machine(resource_group_name=resource_group,
+                                   vm_name=vm_name).result()
+        except Exception as e:  # pylint: disable=broad-except
+            logger.warning('Failed to delete VM: {}'.format(e))
+
+    for nic_name in filtered_resources[_RESOURCE_NETWORK_INTERFACE_TYPE]:
+        try:
+            # Before removing Public IP Address, we need to wait for the
+            # Network Interface to be completely removed with .result() so the
+            # dependency of Network Interface on Public IP Address is
+            # disassociated. This takes about ~1s.
+            _delete_nic_with_retries(network_client, resource_group, nic_name)
+        except Exception as e:  # pylint: disable=broad-except
+            logger.warning('Failed to delete nic: {}'.format(e))
+
+    for public_ip_name in filtered_resources[_RESOURCE_PUBLIC_IP_ADDRESS_TYPE]:
+        try:
+            delete_public_ip_addresses(resource_group_name=resource_group,
+                                       public_ip_address_name=public_ip_name)
+        except Exception as e:  # pylint: disable=broad-except
+            logger.warning('Failed to delete public ip: {}'.format(e))
+
+    for vnet_name in filtered_resources[_RESOURCE_VIRTUAL_NETWORK_TYPE]:
+        try:
+            delete_virtual_networks(resource_group_name=resource_group,
+                                    virtual_network_name=vnet_name)
+        except Exception as e:  # pylint: disable=broad-except
+            logger.warning('Failed to delete vnet: {}'.format(e))
+
+    for msi_name in filtered_resources[_RESOURCE_MANAGED_IDENTITY_TYPE]:
+        user_assigned_identities = (
+            msi_client.user_assigned_identities.list_by_resource_group(
+                resource_group_name=resource_group))
+        for identity in user_assigned_identities:
+            if msi_name == identity.name:
+                # We use the principal_id to find the correct guid converted
+                # role assignment name because each managed identity has a
+                # unique principal_id, and role assignments are associated
+                # with security principals (like managed identities) via this
+                # principal_id.
+                target_principal_id = identity.principal_id
+                scope = (f'/subscriptions/{subscription_id}'
+                         f'/resourceGroups/{resource_group}')
+                role_assignments = auth_client.role_assignments.list_for_scope(
+                    scope)
+                for assignment in role_assignments:
+                    if target_principal_id == assignment.principal_id:
+                        guid_role_assignment_name = assignment.name
+                        try:
+                            delete_role_assignment(
+                                scope=scope,
+                                role_assignment_name=guid_role_assignment_name)
+                        except Exception as e:  # pylint: disable=broad-except
+                            logger.warning('Failed to delete role '
+                                           'assignment: {}'.format(e))
+                        break
+        try:
+            delete_managed_identity(resource_group_name=resource_group,
+                                    resource_name=msi_name)
+        except Exception as e:  # pylint: disable=broad-except
+            logger.warning('Failed to delete msi: {}'.format(e))
+
+    for nsg_name in filtered_resources[_RESOURCE_NETWORK_SECURITY_GROUP_TYPE]:
+        try:
+            delete_network_security_group(resource_group_name=resource_group,
+                                          network_security_group_name=nsg_name)
+        except Exception as e:  # pylint: disable=broad-except
+            logger.warning('Failed to delete nsg: {}'.format(e))
+
+    delete_deployment = _get_azure_sdk_function(
+        client=resource_client.deployments, function_name='begin_delete')
+    deployment_names = [
+        constants.EXTERNAL_RG_BOOTSTRAP_DEPLOYMENT_NAME.format(
+            cluster_name_on_cloud=cluster_name_on_cloud),
+        constants.EXTERNAL_RG_VM_DEPLOYMENT_NAME.format(
+            cluster_name_on_cloud=cluster_name_on_cloud)
+    ]
+    for deployment_name in deployment_names:
+        try:
+            delete_deployment(resource_group_name=resource_group,
+                              deployment_name=deployment_name)
+        except Exception as e:  # pylint: disable=broad-except
+            logger.warning('Failed to delete deployment: {}'.format(e))
+
+
+@common_utils.retry
+def query_instances(
+    cluster_name_on_cloud: str,
+    provider_config: Optional[Dict[str, Any]] = None,
+    non_terminated_only: bool = True,
+) -> Dict[str, Optional[status_lib.ClusterStatus]]:
+    """See sky/provision/__init__.py"""
+    assert provider_config is not None, cluster_name_on_cloud
+
+    subscription_id = provider_config['subscription_id']
+    resource_group = provider_config['resource_group']
+    filters = {constants.TAG_RAY_CLUSTER_NAME: cluster_name_on_cloud}
+    compute_client = azure.get_client('compute', subscription_id)
+    nodes = _filter_instances(compute_client, resource_group, filters)
+    statuses: Dict[str, Optional[status_lib.ClusterStatus]] = {}
+
+    def _fetch_and_map_status(node, resource_group: str) -> None:
+        compute_client = azure.get_client('compute', subscription_id)
+        status = _get_instance_status(compute_client, node, resource_group)
+
+        if status is None and non_terminated_only:
+            return
+        statuses[node.name] = (None if status is None else
+                               status.to_cluster_status())
+
+    with pool.ThreadPool() as p:
+        p.starmap(_fetch_and_map_status,
+                  [(node, resource_group) for node in nodes])
+
+    return statuses
+
+
 def open_ports(
     cluster_name_on_cloud: str,
     ports: List[str],
@@ -45,43 +992,71 @@ def open_ports(
     subscription_id = provider_config['subscription_id']
     resource_group = provider_config['resource_group']
     network_client = azure.get_client('network', subscription_id)
-    # The NSG should have been created by the cluster provisioning.
-    update_network_security_groups = get_azure_sdk_function(
+
+    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(
+    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:
-            # 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':
+        # Given resource group can contain network security groups that are
+        # irrelevant to this provisioning especially with user specified
+        # resource group at ~/.sky/config. So we make sure to check for the
+        # completion of nsg relevant to the VM being provisioned.
+        if cluster_name_on_cloud in nsg.name:
+            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}'
+                            f' in {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':
+                    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:
                 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:
-            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 {nsg.name}.'
+                    ) from e
 
 
 def cleanup_ports(