"""Backend: runs on cloud virtual machines, managed by Ray."""
import ast
import enum
import getpass
import inspect
import math
import json
import os
import pathlib
import re
import signal
import subprocess
import tempfile
import textwrap
import time
import typing
from typing import Dict, List, Optional, Tuple, Union

import colorama
import filelock

import sky
from sky import backends
from sky import clouds
from sky import cloud_stores
from sky import exceptions
from sky import global_user_state
from sky import resources as resources_lib
from sky import sky_logging
from sky import optimizer
from sky import spot as spot_lib
from sky import task as task_lib
from sky.data import data_utils
from sky.data import storage as storage_lib
from sky.backends import backend_utils
from sky.backends import onprem_utils
from sky.backends import wheel_utils
from sky.skylet import autostop_lib
from sky.skylet import constants
from sky.skylet import job_lib
from sky.skylet import log_lib
from sky.usage import usage_lib
from sky.utils import common_utils
from sky.utils import command_runner
from sky.utils import log_utils
from sky.utils import subprocess_utils
from sky.utils import timeline
from sky.utils import tpu_utils
from sky.utils import ux_utils

if typing.TYPE_CHECKING:
    from sky import dag

OptimizeTarget = optimizer.OptimizeTarget
Path = str

SKY_REMOTE_APP_DIR = backend_utils.SKY_REMOTE_APP_DIR
SKY_REMOTE_WORKDIR = constants.SKY_REMOTE_WORKDIR

logger = sky_logging.init_logger(__name__)

_PATH_SIZE_MEGABYTES_WARN_THRESHOLD = 256

# Timeout (seconds) for provision progress: if in this duration no new nodes
# are launched, abort and failover.
_NODES_LAUNCHING_PROGRESS_TIMEOUT = 90

# Time gap between retries after failing to provision in all possible places.
# Used only if --retry-until-up is set.
_RETRY_UNTIL_UP_INIT_GAP_SECONDS = 60

# The maximum retry count for fetching IP address.
_HEAD_IP_MAX_ATTEMPTS = 5
_WORKER_IP_MAX_ATTEMPTS = 5

_TEARDOWN_FAILURE_MESSAGE = (
    f'\n{colorama.Fore.RED}Failed to terminate '
    '{cluster_name}. {extra_reason}'
    'If you want to ignore this error and remove the cluster '
    'from the status table, use `sky down --purge`.'
    f'{colorama.Style.RESET_ALL}\n'
    '**** STDOUT ****\n'
    '{stdout}\n'
    '**** STDERR ****\n'
    '{stderr}')

_TEARDOWN_PURGE_WARNING = (
    f'{colorama.Fore.YELLOW}'
    'WARNING: Received non-zero exit code from {reason}. '
    'Make sure resources are manually deleted.'
    f'{colorama.Style.RESET_ALL}')

_TPU_NOT_FOUND_ERROR = 'ERROR: (gcloud.compute.tpus.delete) NOT_FOUND'

_CTRL_C_TIP_MESSAGE = ('INFO: Tip: use Ctrl-C to exit log streaming '
                       '(task will not be killed).')

_MAX_RAY_UP_RETRY = 5

_JOB_ID_PATTERN = re.compile(r'Job ID: ([0-9]+)')


def _get_cluster_config_template(cloud):
    cloud_to_template = {
        clouds.AWS: 'aws-ray.yml.j2',
        clouds.Azure: 'azure-ray.yml.j2',
        clouds.GCP: 'gcp-ray.yml.j2',
        clouds.Local: 'local-ray.yml.j2',
    }
    return cloud_to_template[type(cloud)]


class RayCodeGen:
    """Code generator of a Ray program that executes a sky.Task.

    Usage:

      >> codegen = RayCodegen()
      >> codegen.add_prologue()

      >> codegen.add_ray_task(...)
      >> codegen.add_ray_task(...)

      >> codegen.add_epilogue()
      >> code = codegen.build()
    """

    def __init__(self):
        # Code generated so far, to be joined via '\n'.
        self._code = []
        # Guard method calling order.
        self._has_prologue = False
        self._has_epilogue = False

        # For n nodes gang scheduling.
        self._has_gang_scheduling = False
        self._num_nodes = 0

        self._has_register_run_fn = False

        # job_id
        # Job ID is used to identify the job (also this generated code).
        # It is a int automatically generated by the DB on the cluster
        # and monotonically increasing starting from 1.
        # To generate the job ID, we use the following logic:
        #   code = job_lib.JobLibCodeGen.add_job(username,
        #                                              run_timestamp)
        #   job_id = get_output(run_on_cluster(code))
        self.job_id = None

    def add_prologue(self,
                     job_id: int,
                     spot_task: Optional['task_lib.Task'] = None,
                     setup_cmd: Optional[str] = None,
                     envs: Optional[Dict[str, str]] = None,
                     setup_log_path: Optional[str] = None,
                     is_local: bool = False) -> None:
        assert not self._has_prologue, 'add_prologue() called twice?'
        self._has_prologue = True
        self.job_id = job_id
        # Should use 'auto' or 'ray://<internal_head_ip>:10001' rather than
        # 'ray://localhost:10001', or 'ray://127.0.0.1:10001', for public cloud.
        # Otherwise, it will a bug of ray job failed to get the placement group
        # in ray <= 2.0.1.
        # TODO(mluo): Check why 'auto' not working with on-prem cluster and
        # whether the placement group issue also occurs in on-prem cluster.
        ray_address = 'ray://localhost:10001' if is_local else 'auto'
        self._code = [
            textwrap.dedent(f"""\
            import getpass
            import hashlib
            import io
            import os
            import pathlib
            import sys
            import selectors
            import subprocess
            import tempfile
            import textwrap
            import time
            from typing import Dict, List, Optional, Tuple, Union

            import ray
            import ray.util as ray_util

            from sky.skylet import constants
            from sky.skylet import job_lib
            from sky.utils import log_utils

            SKY_REMOTE_WORKDIR = {constants.SKY_REMOTE_WORKDIR!r}

            ray.init(address={ray_address!r}, namespace='__sky__{job_id}__', log_to_driver=True)
            run_fn = None
            futures = []
            """),
            # FIXME: This is a hack to make sure that the functions can be found
            # by ray.remote. This should be removed once we have a better way to
            # specify dependencies for ray.
            inspect.getsource(log_lib.process_subprocess_stream),
            inspect.getsource(log_lib.run_with_log),
            inspect.getsource(log_lib.make_task_bash_script),
            inspect.getsource(log_lib.add_ray_env_vars),
            inspect.getsource(log_lib.run_bash_command_with_log),
            'run_bash_command_with_log = ray.remote(run_bash_command_with_log)',
            f'setup_cmd = {setup_cmd!r}'
        ]
        if setup_cmd is not None:
            self._code += [
                textwrap.dedent(f"""\
                _SETUP_CPUS = 0.0001
                # The setup command will be run as a ray task with num_cpus=_SETUP_CPUS as the
                # requirement; this means Ray will set CUDA_VISIBLE_DEVICES to an empty string.
                # We unset it so that user setup command may properly use this env var.
                setup_cmd = 'unset CUDA_VISIBLE_DEVICES; ' + setup_cmd
                job_lib.set_status({job_id!r}, job_lib.JobStatus.SETTING_UP)
                print({_CTRL_C_TIP_MESSAGE!r}, file=sys.stderr, flush=True)
                total_num_nodes = len(ray.nodes())
                setup_bundles = [{{"CPU": _SETUP_CPUS}} for _ in range(total_num_nodes)]
                setup_pg = ray.util.placement_group(setup_bundles, strategy='STRICT_SPREAD')
                ray.get(setup_pg.ready())
                setup_workers = [run_bash_command_with_log \\
                    .options(name='setup', num_cpus=_SETUP_CPUS, placement_group=setup_pg, placement_group_bundle_index=i) \\
                    .remote(
                        setup_cmd,
                        os.path.expanduser({setup_log_path!r}),
                        getpass.getuser(),
                        job_id={self.job_id},
                        env_vars={envs!r},
                        stream_logs=True,
                        with_ray=True,
                        use_sudo={is_local},
                    ) for i in range(total_num_nodes)]
                setup_returncodes = ray.get(setup_workers)
                if sum(setup_returncodes) != 0:
                    job_lib.set_status({self.job_id!r}, job_lib.JobStatus.FAILED_SETUP)
                    # This waits for all streaming logs to finish.
                    time.sleep(1)
                    print('ERROR: {colorama.Fore.RED}Job {self.job_id}\\'s setup failed with '
                        'return code list:{colorama.Style.RESET_ALL}',
                        setup_returncodes,
                        file=sys.stderr,
                        flush=True)
                    # Need this to set the job status in ray job to be FAILED.
                    sys.exit(1)
                """)
            ]
        self._code += [
            f'job_lib.set_status({job_id!r}, job_lib.JobStatus.PENDING)',
        ]
        if spot_task is not None:
            # Add the spot job to spot queue table.
            resources_str = backend_utils.get_task_resources_str(spot_task)
            self._code += [
                'from sky.spot import spot_state',
                f'spot_state.set_pending('
                f'{job_id}, {spot_task.name!r}, {resources_str!r})',
            ]

    def add_gang_scheduling_placement_group(
        self,
        num_nodes: int,
        accelerator_dict: Dict[str, int],
        cluster_ips_sorted: Optional[List[str]] = None,
    ) -> None:
        """Create the gang scheduling placement group for a Task.

        cluster_ips_sorted is used to ensure that the SKY_NODE_RANK environment
        variable is assigned in a deterministic order whenever a new task is
        added.
        """
        assert self._has_prologue, ('Call add_prologue() before '
                                    'add_gang_scheduling_placement_group().')
        self._has_gang_scheduling = True
        self._num_nodes = num_nodes

        # Set CPU to avoid ray hanging the resources allocation
        # for remote functions, since the task will request 1 CPU
        # by default.
        bundles = [{
            'CPU': backend_utils.DEFAULT_TASK_CPU_DEMAND
        } for _ in range(num_nodes)]

        if accelerator_dict is not None:
            acc_name = list(accelerator_dict.keys())[0]
            acc_count = list(accelerator_dict.values())[0]
            gpu_dict = {'GPU': acc_count}
            # gpu_dict should be empty when the accelerator is not GPU.
            # FIXME: This is a hack to make sure that we do not reserve
            # GPU when requesting TPU.
            if 'tpu' in acc_name.lower():
                gpu_dict = dict()
            for bundle in bundles:
                bundle.update({
                    **accelerator_dict,
                    # Set the GPU to avoid ray hanging the resources allocation
                    **gpu_dict,
                })

        self._code += [
            textwrap.dedent(f"""\
                pg = ray_util.placement_group({json.dumps(bundles)}, 'STRICT_SPREAD')
                plural = 's' if {num_nodes} > 1 else ''
                node_str = f'{num_nodes} node{{plural}}'
                
                message = '' if setup_cmd is not None else {_CTRL_C_TIP_MESSAGE!r} + '\\n'
                message += f'INFO: Waiting for task resources on {{node_str}}. This will block if the cluster is full.'
                print(message,
                      file=sys.stderr,
                      flush=True)
                # FIXME: This will print the error message from autoscaler if
                # it is waiting for other task to finish. We should hide the
                # error message.
                ray.get(pg.ready())
                print('INFO: All task resources reserved.',
                      file=sys.stderr,
                      flush=True)
                job_lib.set_job_started({self.job_id!r})
                """)
        ]

        # Export IP and node rank to the environment variables.
        self._code += [
            textwrap.dedent(f"""\
                @ray.remote
                def check_ip():
                    return ray.util.get_node_ip_address()
                gang_scheduling_id_to_ip = ray.get([
                    check_ip.options(num_cpus={backend_utils.DEFAULT_TASK_CPU_DEMAND},
                                     placement_group=pg,
                                     placement_group_bundle_index=i).remote()
                    for i in range(pg.bundle_count)
                ])
                print('INFO: Reserved IPs:', gang_scheduling_id_to_ip)

                if {cluster_ips_sorted!r} is not None:
                    cluster_ips_map = {{ip: i for i, ip in enumerate({cluster_ips_sorted!r})}}
                    ip_rank_list = sorted(gang_scheduling_id_to_ip, key=cluster_ips_map.get)
                    ip_rank_map = {{ip: i for i, ip in enumerate(ip_rank_list)}}
                    ip_list_str = '\\n'.join(ip_rank_list)
                else:
                    ip_rank_map = {{ip: i for i, ip in enumerate(gang_scheduling_id_to_ip)}}
                    ip_list_str = '\\n'.join(gang_scheduling_id_to_ip)
                """),
        ]

    def register_run_fn(self, run_fn: str, run_fn_name: str) -> None:
        """Register the run function to be run on the remote cluster.

        Args:
            run_fn: The run function to be run on the remote cluster.
        """
        assert self._has_gang_scheduling, (
            'Call add_gang_scheduling_placement_group() '
            'before register_run_fn().')
        assert not self._has_register_run_fn, (
            'register_run_fn() called twice?')
        self._has_register_run_fn = True

        self._code += [
            run_fn,
            f'run_fn = {run_fn_name}',
        ]

    def add_ray_task(self,
                     bash_script: str,
                     task_name: Optional[str],
                     job_run_id: Optional[str],
                     ray_resources_dict: Optional[Dict[str, float]],
                     log_path: str,
                     env_vars: Dict[str, str] = None,
                     gang_scheduling_id: int = 0,
                     use_sudo: bool = False) -> None:
        """Generates code for a ray remote task that runs a bash command."""
        assert self._has_gang_scheduling, (
            'Call add_gang_schedule_placement_group() before add_ray_task().')
        assert (not self._has_register_run_fn or
                bash_script is None), ('bash_script should '
                                       'be None when run_fn is registered.')
        # Build remote_task.options(...)
        #   name=...
        #   resources=...
        #   num_gpus=...
        name_str = f'name=\'{task_name}\''
        if task_name is None:
            # Make the task name more meaningful in ray log.
            name_str = 'name=\'task\''
        cpu_str = f', num_cpus={backend_utils.DEFAULT_TASK_CPU_DEMAND}'

        resources_str = ''
        num_gpus = 0
        num_gpus_str = ''
        if ray_resources_dict is not None:
            assert len(ray_resources_dict) == 1, \
                ('There can only be one type of accelerator per instance.'
                 f' Found: {ray_resources_dict}.')
            num_gpus = list(ray_resources_dict.values())[0]
            resources_str = f', resources={json.dumps(ray_resources_dict)}'

            # Passing this ensures that the Ray remote task gets
            # CUDA_VISIBLE_DEVICES set correctly.  If not passed, that flag
            # would be force-set to empty by Ray.
            num_gpus_str = f', num_gpus={num_gpus}'
            # `num_gpus` should be empty when the accelerator is not GPU.
            # FIXME: use a set of GPU types.
            resources_key = list(ray_resources_dict.keys())[0]
            if 'tpu' in resources_key.lower():
                num_gpus_str = ''
        resources_str += ', placement_group=pg'
        resources_str += f', placement_group_bundle_index={gang_scheduling_id}'

        sky_env_vars_dict_str = [
            textwrap.dedent("""\
            sky_env_vars_dict = dict()
            sky_env_vars_dict['SKYPILOT_NODE_IPS'] = ip_list_str
            # Environment starting with `SKY_` is deprecated.
            sky_env_vars_dict['SKY_NODE_IPS'] = ip_list_str
            """)
        ]

        if env_vars is not None:
            sky_env_vars_dict_str.extend(f'sky_env_vars_dict[{k!r}] = {v!r}'
                                         for k, v in env_vars.items())
        if job_run_id is not None:
            sky_env_vars_dict_str += [
                f'sky_env_vars_dict[{constants.JOB_ID_ENV_VAR!r}]'
                f' = {job_run_id!r}'
            ]
        sky_env_vars_dict_str = '\n'.join(sky_env_vars_dict_str)

        logger.debug('Added Task with options: '
                     f'{name_str}{cpu_str}{resources_str}{num_gpus_str}')
        self._code += [
            sky_env_vars_dict_str,
            textwrap.dedent(f"""\
        script = {bash_script!r}
        if run_fn is not None:
            script = run_fn({gang_scheduling_id}, gang_scheduling_id_to_ip)

        log_path = os.path.expanduser({log_path!r})

        if script is not None:
            sky_env_vars_dict['SKYPILOT_NUM_GPUS_PER_NODE'] = {int(math.ceil(num_gpus))!r}
            # Environment starting with `SKY_` is deprecated.
            sky_env_vars_dict['SKY_NUM_GPUS_PER_NODE'] = {int(math.ceil(num_gpus))!r}

            ip = gang_scheduling_id_to_ip[{gang_scheduling_id!r}]
            sky_env_vars_dict['SKYPILOT_NODE_RANK'] = ip_rank_map[ip]
            # Environment starting with `SKY_` is deprecated.
            sky_env_vars_dict['SKY_NODE_RANK'] = ip_rank_map[ip]
            
            sky_env_vars_dict['SKYPILOT_INTERNAL_JOB_ID'] = {self.job_id}
            # Environment starting with `SKY_` is deprecated.
            sky_env_vars_dict['SKY_INTERNAL_JOB_ID'] = {self.job_id}

            futures.append(run_bash_command_with_log \\
                    .options({name_str}{cpu_str}{resources_str}{num_gpus_str}) \\
                    .remote(
                        script,
                        log_path,
                        getpass.getuser(),
                        job_id={self.job_id},
                        env_vars=sky_env_vars_dict,
                        stream_logs=True,
                        with_ray=True,
                        use_sudo={use_sudo},
                    ))""")
        ]

    def add_epilogue(self) -> None:
        """Generates code that waits for all tasks, then exits."""
        assert self._has_prologue, 'Call add_prologue() before add_epilogue().'
        assert not self._has_epilogue, 'add_epilogue() called twice?'
        self._has_epilogue = True

        self._code += [
            textwrap.dedent(f"""\
            returncodes = ray.get(futures)
            if sum(returncodes) != 0:
                job_lib.set_status({self.job_id!r}, job_lib.JobStatus.FAILED)
                # This waits for all streaming logs to finish.
                time.sleep(1)
                print('ERROR: {colorama.Fore.RED}Job {self.job_id} failed with '
                      'return code list:{colorama.Style.RESET_ALL}',
                      returncodes,
                      file=sys.stderr,
                      flush=True)
                # Need this to set the job status in ray job to be FAILED.
                sys.exit(1)
            else:
                sys.stdout.flush()
                sys.stderr.flush()
                job_lib.set_status({self.job_id!r}, job_lib.JobStatus.SUCCEEDED)
                # This waits for all streaming logs to finish.
                time.sleep(1)
            """)
        ]

    def build(self) -> str:
        """Returns the entire generated program."""
        assert self._has_epilogue, 'Call add_epilogue() before build().'
        return '\n'.join(self._code)


class RetryingVmProvisioner(object):
    """A provisioner that retries different cloud/regions/zones."""

    class ToProvisionConfig:
        """Resources to be provisioned."""

        def __init__(self,
                     cluster_name: str,
                     resources: Optional[resources_lib.Resources],
                     num_nodes: int,
                     cluster_exists: bool = False) -> None:
            assert cluster_name is not None, 'cluster_name must be specified.'
            self.cluster_name = cluster_name
            self.resources = resources
            self.num_nodes = num_nodes
            # Whether the cluster exists in the clusters table. It may be
            # actually live or down.
            self.cluster_exists = cluster_exists

    class GangSchedulingStatus(enum.Enum):
        """Enum for gang scheduling status."""
        CLUSTER_READY = 0
        GANG_FAILED = 1
        HEAD_FAILED = 2

    def __init__(self, log_dir: str, dag: 'dag.Dag',
                 optimize_target: OptimizeTarget,
                 local_wheel_path: pathlib.Path, wheel_hash: str):
        self._blocked_regions = set()
        self._blocked_zones = set()
        self._blocked_launchable_resources = set()

        self.log_dir = os.path.expanduser(log_dir)
        self._dag = dag
        self._optimize_target = optimize_target
        self._local_wheel_path = local_wheel_path
        self._wheel_hash = wheel_hash

        colorama.init()

    def _in_blocklist(self, cloud, region, zones):
        if region.name in self._blocked_regions:
            return True
        # We do not keep track of zones in Azure and Local,
        # as both clouds do not have zones.
        if isinstance(cloud, (clouds.Azure, clouds.Local)):
            return False
        assert zones, (cloud, region, zones)
        for zone in zones:
            if zone.name not in self._blocked_zones:
                return False
        return True

    def _clear_blocklist(self):
        self._blocked_regions.clear()
        self._blocked_zones.clear()

    def _update_blocklist_on_gcp_error(self, region, zones, stdout, stderr):
        style = colorama.Style
        assert len(zones) == 1, zones
        zone = zones[0]
        splits = stderr.split('\n')
        exception_str = [s for s in splits if s.startswith('Exception: ')]
        httperror_str = [
            s for s in splits
            if s.startswith('googleapiclient.errors.HttpError: ')
        ]
        if len(exception_str) == 1:
            # Parse structured response {'errors': [...]}.
            exception_str = exception_str[0][len('Exception: '):]
            try:
                exception_dict = ast.literal_eval(exception_str)
            except Exception as e:
                raise RuntimeError(
                    f'Failed to parse exception: {exception_str}') from e
            # TPU VM returns a different structured response.
            if 'errors' not in exception_dict:
                exception_dict = {'errors': [exception_dict]}
            for error in exception_dict['errors']:
                code = error['code']
                message = error['message']
                logger.warning(f'Got return code {code} in {zone.name} '
                               f'{style.DIM}(message: {message})'
                               f'{style.RESET_ALL}')
                if code == 'QUOTA_EXCEEDED':
                    if '\'GPUS_ALL_REGIONS\' exceeded' in message:
                        # Global quota.  All regions in GCP will fail.  Ex:
                        # Quota 'GPUS_ALL_REGIONS' exceeded.  Limit: 1.0
                        # globally.
                        # This skip is only correct if we implement "first
                        # retry the region/zone of an existing cluster with the
                        # same name" correctly.
                        for r, _ in clouds.GCP.region_zones_provision_loop():
                            self._blocked_regions.add(r.name)
                    else:
                        # Per region.  Ex: Quota 'CPUS' exceeded.  Limit: 24.0
                        # in region us-west1.
                        self._blocked_regions.add(region.name)
                elif code in [
                        'ZONE_RESOURCE_POOL_EXHAUSTED',
                        'ZONE_RESOURCE_POOL_EXHAUSTED_WITH_DETAILS',
                        'UNSUPPORTED_OPERATION'
                ]:  # Per zone.
                    # Return codes can be found at https://cloud.google.com/compute/docs/troubleshooting/troubleshooting-vm-creation # pylint: disable=line-too-long
                    # However, UNSUPPORTED_OPERATION is observed empirically
                    # when VM is preempted during creation.  This seems to be
                    # not documented by GCP.
                    self._blocked_zones.add(zone.name)
                elif code in ['RESOURCE_NOT_READY']:
                    # This code is returned when the VM is still STOPPING.
                    self._blocked_zones.add(zone.name)
                elif code == 8:
                    # Error code 8 means TPU resources is out of
                    # capacity. Example:
                    # {'code': 8, 'message': 'There is no more capacity in the zone "europe-west4-a"; you can try in another zone where Cloud TPU Nodes are offered (see https://cloud.google.com/tpu/docs/regions) [EID: 0x1bc8f9d790be9142]'} # pylint: disable=line-too-long
                    self._blocked_zones.add(zone.name)
                else:
                    assert False, error
        elif len(httperror_str) >= 1:
            # Parse HttpError for unauthorized regions. Example:
            # googleapiclient.errors.HttpError: <HttpError 403 when requesting ... returned "Location us-east1-d is not found or access is unauthorized.". # pylint: disable=line-too-long
            # Details: "Location us-east1-d is not found or access is
            # unauthorized.">
            logger.info(f'Got {httperror_str[0]}')
            self._blocked_zones.add(zone.name)
        else:
            # No such structured error response found.
            assert not exception_str, stderr
            if 'was not found' in stderr:
                # Example: The resource
                # 'projects/<id>/zones/zone/acceleratorTypes/nvidia-tesla-v100'
                # was not found.
                logger.warning(f'Got \'resource not found\' in {zone.name}.')
                self._blocked_zones.add(zone.name)
            else:
                logger.info('====== stdout ======')
                for s in stdout.split('\n'):
                    print(s)
                logger.info('====== stderr ======')
                for s in splits:
                    print(s)

                with ux_utils.print_exception_no_traceback():
                    raise RuntimeError('Errors occurred during provision; '
                                       'check logs above.')

    def _update_blocklist_on_aws_error(self, region, zones, stdout, stderr):
        del zones  # Unused.
        style = colorama.Style
        stdout_splits = stdout.split('\n')
        stderr_splits = stderr.split('\n')
        errors = [
            s.strip()
            for s in stdout_splits + stderr_splits
            if 'An error occurred' in s.strip()
        ]
        # Need to handle boto3 printing error but its retry succeeded:
        #   error occurred (Unsupported) .. not supported in your requested
        #   Availability Zone (us-west-2d)...retrying
        #   --> it automatically succeeded in another zone
        #   --> failed in [4/7] Running initialization commands due to user cmd
        # In this case, we should error out.
        head_node_up = any(
            line.startswith('<1/1> Setting up head node')
            for line in stdout_splits + stderr_splits)
        if not errors or head_node_up:
            # TODO: Got transient 'Failed to create security group' that goes
            # away after a few minutes.  Should we auto retry other regions, or
            # let the user retry.
            logger.info('====== stdout ======')
            for s in stdout_splits:
                print(s)
            logger.info('====== stderr ======')
            for s in stderr_splits:
                print(s)
            with ux_utils.print_exception_no_traceback():
                raise RuntimeError('Errors occurred during provision; '
                                   'check logs above.')
        # The underlying ray autoscaler / boto3 will try all zones of a region
        # at once.
        logger.warning(f'Got error(s) in all zones of {region.name}:')
        messages = '\n\t'.join(errors)
        logger.warning(f'{style.DIM}\t{messages}{style.RESET_ALL}')
        self._blocked_regions.add(region.name)

    def _update_blocklist_on_azure_error(self, region, zones, stdout, stderr):
        del zones  # Unused.
        # The underlying ray autoscaler will try all zones of a region at once.
        style = colorama.Style
        stdout_splits = stdout.split('\n')
        stderr_splits = stderr.split('\n')
        errors = [
            s.strip()
            for s in stdout_splits + stderr_splits
            if ('Exception Details:' in s.strip() or 'InvalidTemplateDeployment'
                in s.strip() or '(ReadOnlyDisabledSubscription)' in s.strip())
        ]
        if not errors:
            logger.info('====== stdout ======')
            for s in stdout_splits:
                print(s)
            logger.info('====== stderr ======')
            for s in stderr_splits:
                print(s)
            with ux_utils.print_exception_no_traceback():
                raise RuntimeError('Errors occurred during provision; '
                                   'check logs above.')

        logger.warning(f'Got error(s) in {region.name}:')
        messages = '\n\t'.join(errors)
        logger.warning(f'{style.DIM}\t{messages}{style.RESET_ALL}')
        if any('(ReadOnlyDisabledSubscription)' in s for s in errors):
            for r in sky.Azure.regions():
                self._blocked_regions.add(r.name)
        else:
            self._blocked_regions.add(region.name)

    def _update_blocklist_on_local_error(self, region, zones, stdout, stderr):
        del zones  # Unused.
        style = colorama.Style
        stdout_splits = stdout.split('\n')
        stderr_splits = stderr.split('\n')
        errors = [
            s.strip()
            for s in stdout_splits + stderr_splits
            if 'ERR' in s.strip() or 'PANIC' in s.strip()
        ]
        if not errors:
            logger.info('====== stdout ======')
            for s in stdout_splits:
                print(s)
            logger.info('====== stderr ======')
            for s in stderr_splits:
                print(s)
            with ux_utils.print_exception_no_traceback():
                raise RuntimeError(
                    'Errors occurred during launching of cluster services; '
                    'check logs above.')
        logger.warning('Got error(s) on local cluster:')
        messages = '\n\t'.join(errors)
        logger.warning(f'{style.DIM}\t{messages}{style.RESET_ALL}')
        self._blocked_regions.add(region.name)

    def _update_blocklist_on_error(self, cloud, region, zones, stdout,
                                   stderr) -> None:
        """Handles cloud-specific errors and updates the block list.

        This parses textual stdout/stderr because we don't directly use the
        underlying clouds' SDKs.  If we did that, we could catch proper
        exceptions instead.
        """
        if stdout is None:
            # Gang scheduling failure.  Simply block the region.
            assert stderr is None, stderr
            self._blocked_regions.add(region.name)
            return

        if isinstance(cloud, clouds.GCP):
            return self._update_blocklist_on_gcp_error(region, zones, stdout,
                                                       stderr)

        if isinstance(cloud, clouds.AWS):
            return self._update_blocklist_on_aws_error(region, zones, stdout,
                                                       stderr)

        if isinstance(cloud, clouds.Azure):
            return self._update_blocklist_on_azure_error(
                region, zones, stdout, stderr)

        if isinstance(cloud, clouds.Local):
            return self._update_blocklist_on_local_error(
                region, zones, stdout, stderr)

        assert False, f'Unknown cloud: {cloud}.'

    def _yield_region_zones(self, to_provision: resources_lib.Resources,
                            cluster_name: str, cluster_exists: bool):
        cloud = to_provision.cloud
        region = None
        zones = None
        # Try loading previously launched region/zones and try them first,
        # because we may have an existing cluster there.
        # Get the *previous* cluster status and handle.
        cluster_status, handle = backend_utils.refresh_cluster_status_handle(
            cluster_name, acquire_per_cluster_status_lock=False)
        if handle is not None:
            try:
                config = common_utils.read_yaml(handle.cluster_yaml)
                prev_resources = handle.launched_resources
                if prev_resources is not None and cloud.is_same_cloud(
                        prev_resources.cloud):
                    if cloud.is_same_cloud(clouds.GCP()) or cloud.is_same_cloud(
                            clouds.AWS()):
                        region = config['provider']['region']
                        zones = config['provider']['availability_zone']
                    elif cloud.is_same_cloud(clouds.Azure()):
                        region = config['provider']['location']
                        zones = None
                    elif cloud.is_same_cloud(clouds.Local()):
                        local_regions = clouds.Local.regions()
                        region = local_regions[0].name
                        zones = None
                    else:
                        assert False, cloud
                    assert region == prev_resources.region, (
                        f'Region mismatch. The region in '
                        f'{handle.cluster_yaml} '
                        'has been changed from '
                        f'{prev_resources.region} to {region}.')
                    assert (zones is None or prev_resources.zone is None or
                            prev_resources.zone
                            in zones), (f'{prev_resources.zone} not found in '
                                        f'zones of {handle.cluster_yaml}.')
                    # Note that we don't overwrite the zone field in Ray YAML
                    # even if prev_resources.zone != zones.
                    # This is because Ray will consider the YAML hash changed
                    # and not reuse the existing cluster.
            except FileNotFoundError:
                # Happens if no previous cluster.yaml exists.
                pass
        if region is not None and cluster_exists:

            region = clouds.Region(name=region)
            if zones is not None:
                zones = [clouds.Zone(name=zone) for zone in zones.split(',')]
                region.set_zones(zones)

            if cluster_status != global_user_state.ClusterStatus.UP:
                logger.info(
                    f'Cluster {cluster_name!r} (status: {cluster_status.value})'
                    f' was previously launched in {cloud} '
                    f'({region.name}). Relaunching in that region.')
            # This should be handled in the check_resources_match function.
            assert (to_provision.region is None or
                    region.name == to_provision.region), (
                        f'Cluster {cluster_name!r} was previously launched in '
                        f'{cloud} ({region.name}). It does not match the '
                        f'required region {to_provision.region}.')
            # TODO(zhwu): The cluster being killed by cloud provider should
            # be tested whether re-launching a cluster killed spot instance
            # will recover the data.
            yield (region, zones)  # Ok to yield again in the next loop.

            # Cluster with status UP can reach here, if it was killed by the
            # cloud provider and no available resources in that region to
            # relaunch, which can happen to spot instance.
            if cluster_status == global_user_state.ClusterStatus.UP:
                message = (
                    f'Failed to connect to the cluster {cluster_name!r}. '
                    'It is possibly killed by cloud provider or manually '
                    'in the cloud provider console. To remove the cluster '
                    f'please run: sky down {cluster_name}')
                # Reset to UP (rather than keeping it at INIT), as INIT
                # mode will enable failover to other regions, causing
                # data lose.
                # TODO(zhwu): This is set to UP to be more conservative,
                # we may need to confirm whether the cluster is down in all
                # cases.
                global_user_state.set_cluster_status(
                    cluster_name, global_user_state.ClusterStatus.UP)
                with ux_utils.print_exception_no_traceback():
                    raise exceptions.ResourcesUnavailableError(message,
                                                               no_failover=True)

            # If it reaches here: the cluster status gets set to INIT, since
            # a launch request was issued but failed.
            #
            # Check the *previous* cluster status. If the cluster is previously
            # stopped, we should not retry other regions, since the previously
            # attached volumes are not visible on another region.
            if cluster_status == global_user_state.ClusterStatus.STOPPED:
                message = (
                    'Failed to acquire resources to restart the stopped '
                    f'cluster {cluster_name} on {region}. Please retry again '
                    'later.')

                # Reset to STOPPED (rather than keeping it at INIT), because
                # (1) the cluster is not up (2) it ensures future `sky start`
                # will disable auto-failover too.
                global_user_state.set_cluster_status(
                    cluster_name, global_user_state.ClusterStatus.STOPPED)

                with ux_utils.print_exception_no_traceback():
                    raise exceptions.ResourcesUnavailableError(message,
                                                               no_failover=True)

            assert cluster_status == global_user_state.ClusterStatus.INIT
            message = (f'Failed to launch cluster {cluster_name!r} '
                       f'(previous status: {cluster_status.value}) '
                       f'with the original resources: {to_provision}.')
            # We attempted re-launching a previously INIT cluster with the
            # same cloud/region/resources, but failed. Here no_failover=False,
            # so we will retry provisioning it with the current requested
            # resources in the outer loop.
            #
            # This condition can be triggered for previously INIT cluster by
            # (1) launch, after answering prompt immediately ctrl-c;
            # (2) launch again.
            # After (1), the cluster exists with INIT, and may or may not be
            # live.  And if it hits here, it's definitely not alive (because
            # step (2) failed).  Hence it's ok to retry with different
            # cloud/region and with current resources.
            with ux_utils.print_exception_no_traceback():
                raise exceptions.ResourcesUnavailableError(message)

        for region, zones in cloud.region_zones_provision_loop(
                instance_type=to_provision.instance_type,
                accelerators=to_provision.accelerators,
                use_spot=to_provision.use_spot,
        ):
            # Only retry requested region/zones or all if not specified.
            zone_names = [zone.name for zone in zones]
            if not to_provision.valid_on_region_zones(region.name, zone_names):
                continue
            if to_provision.zone is not None:
                zones = [clouds.Zone(name=to_provision.zone)]
            yield (region, zones)

    def _try_provision_tpu(self, to_provision: resources_lib.Resources,
                           config_dict: Dict[str, str]) -> bool:
        """Returns whether the provision is successful."""
        tpu_name = config_dict['tpu_name']
        assert 'tpu-create-script' in config_dict, \
            'Expect TPU provisioning with gcloud.'
        try:
            with backend_utils.safe_console_status(
                    '[bold cyan]Provisioning TPU '
                    f'[green]{tpu_name}[/]'):
                subprocess_utils.run(f'bash {config_dict["tpu-create-script"]}',
                                     stdout=subprocess.PIPE,
                                     stderr=subprocess.PIPE)
            return True
        except subprocess.CalledProcessError as e:
            stderr = e.stderr.decode('ascii')
            if 'ALREADY_EXISTS' in stderr:
                # FIXME: should use 'start' on stopped TPUs, replacing
                # 'create'. Or it can be in a "deleting" state. Investigate the
                # right thing to do (force kill + re-provision?).
                logger.info(
                    f'  TPU {tpu_name} already exists; skipped creation.')
                return True

            if 'RESOURCE_EXHAUSTED' in stderr:
                with ux_utils.print_exception_no_traceback():
                    raise exceptions.ResourcesUnavailableError(
                        f'TPU {tpu_name} creation failed due to quota '
                        'exhaustion. Please visit '
                        'https://console.cloud.google.com/iam-admin/quotas '
                        'for more information.')

            if 'PERMISSION_DENIED' in stderr:
                logger.info('  TPUs are not available in this zone.')
                return False

            if 'no more capacity in the zone' in stderr:
                logger.info('  No more capacity in this zone.')
                return False

            if 'CloudTpu received an invalid AcceleratorType' in stderr:
                # INVALID_ARGUMENT: CloudTpu received an invalid
                # AcceleratorType, "v3-8" for zone "us-central1-c". Valid
                # values are "v2-8, ".
                tpu_type = list(to_provision.accelerators.keys())[0]
                logger.info(
                    f'  TPU type {tpu_type} is not available in this zone.')
                return False

            logger.error(stderr)
            raise e

    def _retry_region_zones(self,
                            to_provision: resources_lib.Resources,
                            num_nodes: int,
                            dryrun: bool,
                            stream_logs: bool,
                            cluster_name: str,
                            cluster_exists: bool = False):
        """The provision retry loop."""
        style = colorama.Style
        fore = colorama.Fore
        # Get log_path name
        log_path = os.path.join(self.log_dir, 'provision.log')
        log_abs_path = os.path.abspath(log_path)
        tail_cmd = f'tail -n100 -f {log_path}'
        logger.info('To view detailed progress: '
                    f'{style.BRIGHT}{tail_cmd}{style.RESET_ALL}')

        # Get previous cluster status
        prev_cluster_status = backend_utils.refresh_cluster_status_handle(
            cluster_name, acquire_per_cluster_status_lock=False)[0]

        self._clear_blocklist()
        for region, zones in self._yield_region_zones(to_provision,
                                                      cluster_name,
                                                      cluster_exists):
            if self._in_blocklist(to_provision.cloud, region, zones):
                continue
            zone_str = ','.join(
                z.name for z in zones) if zones is not None else 'all zones'
            config_dict = backend_utils.write_cluster_config(
                to_provision,
                num_nodes,
                _get_cluster_config_template(to_provision.cloud),
                cluster_name,
                self._local_wheel_path,
                self._wheel_hash,
                region=region,
                zones=zones,
                dryrun=dryrun,
                keep_launch_fields_in_existing_config=prev_cluster_status
                is not None)
            if dryrun:
                return
            cluster_config_file = config_dict['ray']

            # Record early, so if anything goes wrong, 'sky status' will show
            # the cluster name and users can appropriately 'sky down'.  It also
            # means a second 'sky launch -c <name>' will attempt to reuse.
            handle = CloudVmRayBackend.ResourceHandle(
                cluster_name=cluster_name,
                cluster_yaml=cluster_config_file,
                launched_nodes=num_nodes,
                # OK for this to be shown in CLI as status == INIT.
                launched_resources=to_provision.copy(region=region.name),
                tpu_create_script=config_dict.get('tpu-create-script'),
                tpu_delete_script=config_dict.get('tpu-delete-script'))
            usage_lib.messages.usage.update_final_cluster_status(
                global_user_state.ClusterStatus.INIT)

            # This sets the status to INIT (even for a normal, UP cluster).
            global_user_state.add_or_update_cluster(cluster_name,
                                                    cluster_handle=handle,
                                                    ready=False)

            tpu_name = config_dict.get('tpu_name')
            if tpu_name is not None:
                logger.info(
                    f'{colorama.Style.BRIGHT}Provisioning TPU on '
                    f'{to_provision.cloud} '
                    f'{region.name}{colorama.Style.RESET_ALL} ({zone_str})')

                success = self._try_provision_tpu(to_provision, config_dict)
                if not success:
                    continue

            logging_info = {
                'cluster_name': cluster_name,
                'region_name': region.name,
                'zone_str': zone_str,
            }
            status, stdout, stderr, head_ip = self._gang_schedule_ray_up(
                to_provision.cloud, num_nodes, cluster_config_file, handle,
                log_abs_path, stream_logs, logging_info, to_provision.use_spot)

            if status == self.GangSchedulingStatus.CLUSTER_READY:
                if cluster_exists:
                    # Guard against the case where there's an existing cluster
                    # with ray runtime messed up (e.g., manually killed) by (1)
                    # querying ray status (2) restarting ray if needed.
                    #
                    # The above 'ray up' will not restart it automatically due
                    # to 'ray up # --no-restart' flag.
                    #
                    # NOTE: this is performance sensitive and has been observed
                    # to take 9s. Only do this for existing clusters, not
                    # freshly launched ones (which should have ray runtime
                    # started).
                    self._ensure_cluster_ray_started(handle, log_abs_path)

                cluster_name = config_dict['cluster_name']
                config_dict['launched_resources'] = to_provision.copy(
                    region=region.name)
                config_dict['launched_nodes'] = num_nodes
                # Optimizations: head_ip and zones may or may not be None. In
                # the latter case, the caller doesn't need to query them again.
                config_dict['head_ip'] = head_ip
                config_dict['zones'] = zones
                plural = '' if num_nodes == 1 else 's'
                if not isinstance(to_provision.cloud, clouds.Local):
                    logger.info(f'{fore.GREEN}Successfully provisioned or found'
                                f' existing VM{plural}.{style.RESET_ALL}')
                return config_dict

            # The cluster is not ready.

            # If cluster was previously UP or STOPPED, stop it; otherwise
            # terminate.
            # FIXME(zongheng): terminating a potentially live cluster is
            # scary. Say: users have an existing cluster that got into INIT, do
            # sky launch, somehow failed, then we may be terminating it here.
            need_terminate = prev_cluster_status not in [
                global_user_state.ClusterStatus.STOPPED,
                global_user_state.ClusterStatus.UP
            ]
            if status == self.GangSchedulingStatus.HEAD_FAILED:
                # ray up failed for the head node.
                self._update_blocklist_on_error(to_provision.cloud, region,
                                                zones, stdout, stderr)
            else:
                # gang scheduling failed.
                assert status == self.GangSchedulingStatus.GANG_FAILED, status
                # The stdout/stderr of ray up is not useful here, since
                # head node is successfully provisioned.
                self._update_blocklist_on_error(
                    to_provision.cloud,
                    region,
                    # Ignored and block region:
                    zones=None,
                    stdout=None,
                    stderr=None)

                # Only log the errors for GANG_FAILED, since HEAD_FAILED may
                # not have created any resources (it can happen however) and
                # HEAD_FAILED can happen in "normal" failover cases.
                logger.error('*** Failed provisioning the cluster. ***')
                terminate_str = 'Terminating' if need_terminate else 'Stopping'
                logger.error(f'*** {terminate_str} the failed cluster. ***')

            # There may exists partial nodes (e.g., head node) so we must
            # terminate or stop before moving on to other regions.
            #
            # NOTE: even HEAD_FAILED could've left a live head node there, so
            # we must terminate/stop here too. E.g., node is up, and ray
            # autoscaler proceeds to setup commands, which may fail:
            #   ERR updater.py:138 -- New status: update-failed
            CloudVmRayBackend().teardown_no_lock(handle,
                                                 terminate=need_terminate)

        message = ('Failed to acquire resources in all regions/zones of '
                   f'{to_provision.cloud}. '
                   'Try changing resource requirements or use another cloud.')
        raise exceptions.ResourcesUnavailableError(message)

    def _tpu_pod_setup(self, cluster_yaml: str,
                       cluster_handle: 'backends.Backend.ResourceHandle',
                       num_nodes: int):
        """Completes setup and start Ray cluster on TPU VM Pod nodes.

        This is a workaround for Ray Autoscaler where `ray up` does not
        run setup or launch ray cluster on TPU VM Pod nodes.
        """
        ssh_credentials = backend_utils.ssh_credential_from_yaml(cluster_yaml)
        all_ips = backend_utils.get_node_ips(cluster_yaml,
                                             num_nodes,
                                             handle=cluster_handle)
        num_tpu_devices = tpu_utils.get_num_tpu_devices(
            cluster_handle.launched_resources)
        if len(all_ips) != num_tpu_devices:
            raise RuntimeError(
                f'Number of nodes IPs: {len(all_ips)} does not'
                f'match number of TPU devices: {num_tpu_devices}.')

        # Get the private IP of head node for connecting Ray cluster.
        head_runner = command_runner.SSHCommandRunner(all_ips[0],
                                                      **ssh_credentials)
        cmd_str = 'python3 -c \"import ray; print(ray._private.services.get_node_ip_address())\"'  # pylint: disable=line-too-long
        rc, stdout, stderr = head_runner.run(cmd_str,
                                             stream_logs=False,
                                             require_outputs=True)
        subprocess_utils.handle_returncode(
            rc,
            cmd_str,
            'Failed to get private IP from head node.',
            stderr=stdout + stderr)
        head_ip_private = stdout.strip()

        ray_config = common_utils.read_yaml(cluster_yaml)
        worker_start_ray_commands = [f'echo "export RAY_HEAD_IP={head_ip_private}" >> ~/.bashrc && source ~/.bashrc']  # pylint: disable=line-too-long
        worker_start_ray_commands += ray_config['worker_start_ray_commands']

        # Setup TPU VM Pod workers and launch Ray cluster.
        backend_utils.do_filemounts_and_setup_on_local_workers(
            cluster_yaml,
            worker_ips=all_ips[1:],
            extra_setup_cmds=worker_start_ray_commands)

    @timeline.event
    def _gang_schedule_ray_up(
            self, to_provision_cloud: clouds.Cloud, num_nodes: int,
            cluster_config_file: str,
            cluster_handle: 'backends.Backend.ResourceHandle',
            log_abs_path: str, stream_logs: bool, logging_info: dict,
            use_spot: bool
    ) -> Tuple[GangSchedulingStatus, str, str, Optional[str]]:
        """Provisions a cluster via 'ray up' and wait until fully provisioned.

        Returns:
          (GangSchedulingStatus; stdout; stderr; optional head_ip).
        """
        # FIXME(zhwu,zongheng): ray up on multiple nodes ups the head node then
        # waits for all workers; turn it into real gang scheduling.
        # FIXME: refactor code path to remove use of stream_logs
        del stream_logs

        style = colorama.Style

        def ray_up():
            # Redirect stdout/err to the file and streaming (if stream_logs).
            # With stdout/err redirected, 'ray up' will have no color and
            # different order from directly running in the console. The
            # `--log-style` and `--log-color` flags do not work. To reproduce,
            # `ray up --log-style pretty --log-color true | tee tmp.out`.

            returncode, stdout, stderr = log_lib.run_with_log(
                # NOTE: --no-restart solves the following bug.  Without it, if
                # 'ray up' (sky launch) twice on a cluster with >1 node, the
                # worker node gets disconnected/killed by ray autoscaler; the
                # whole task will just freeze.  (Doesn't affect 1-node
                # clusters.)  With this flag, ray processes no longer restart
                # and this bug doesn't show.  Downside is existing tasks on the
                # cluster will keep running (which may be ok with the semantics
                # of 'sky launch' twice).
                # Tracked in https://github.com/ray-project/ray/issues/20402.
                ['ray', 'up', '-y', '--no-restart', cluster_config_file],
                log_abs_path,
                stream_logs=False,
                start_streaming_at='Shared connection to',
                line_processor=log_utils.RayUpLineProcessor(),
                # Reduce BOTO_MAX_RETRIES from 12 to 5 to avoid long hanging
                # time during 'ray up' if insufficient capacity occurs.
                env=dict(
                    os.environ,
                    BOTO_MAX_RETRIES='5',
                    # Use environment variables to disable the ray usage stats
                    # (to avoid the 10 second wait for usage collection
                    # confirmation), as the ray version on the user's machine
                    # may be lower version that does not support the
                    # `--disable-usage-stats` flag.
                    RAY_USAGE_STATS_ENABLED='0'),
                require_outputs=True,
                # Disable stdin to avoid ray outputs mess up the terminal with
                # misaligned output when multithreading/multiprocessing are used
                # Refer to: https://github.com/ray-project/ray/blob/d462172be7c5779abf37609aed08af112a533e1e/python/ray/autoscaler/_private/subprocess_output_util.py#L264 # pylint: disable=line-too-long
                stdin=subprocess.DEVNULL)
            return returncode, stdout, stderr

        region_name = logging_info['region_name']
        zone_str = logging_info['zone_str']

        if isinstance(to_provision_cloud, clouds.Local):
            cluster_name = logging_info['cluster_name']
            logger.info(f'{colorama.Style.BRIGHT}Launching on local cluster '
                        f'{cluster_name!r}.')
        else:
            logger.info(
                f'{colorama.Style.BRIGHT}Launching on {to_provision_cloud} '
                f'{region_name}{colorama.Style.RESET_ALL} ({zone_str})')
        start = time.time()

        # Edge case: /tmp/ray does not exist, so autoscaler can't create/store
        # cluster lock and cluster state.
        os.makedirs('/tmp/ray', exist_ok=True)

        # Launch the cluster with ray up

        # Retry if the any of the following happens:
        # 1. Failed due to timeout when fetching head node for Azure.
        # 2. Failed due to file mounts, because it is probably has too
        # many ssh connections and can be fixed by retrying.
        # This is required when using custom image for GCP.
        def need_ray_up(
                ray_up_return_value: Optional[Tuple[int, str, str]]) -> bool:

            # Indicates the first ray up.
            if ray_up_return_value is None:
                return True

            returncode, stdout, stderr = ray_up_return_value
            if returncode == 0:
                return False

            if isinstance(to_provision_cloud, clouds.Azure):
                if 'Failed to invoke the Azure CLI' in stderr:
                    logger.info(
                        'Retrying head node provisioning due to Azure CLI '
                        'issues.')
                    return True
                if ('Head node fetch timed out. Failed to create head node.'
                        in stderr):
                    logger.info(
                        'Retrying head node provisioning due to head fetching '
                        'timeout.')
                    return True

            if isinstance(to_provision_cloud, clouds.GCP):
                if ('Quota exceeded for quota metric \'List requests\' and '
                        'limit \'List requests per minute\'' in stderr):
                    logger.info(
                        'Retrying due to list request rate limit exceeded.')
                    return True

            if ('Processing file mounts' in stdout and
                    'Running setup commands' not in stdout and
                    'Failed to setup head node.' in stderr):
                logger.info(
                    'Retrying runtime setup due to ssh connection issue.')
                return True

            if ('ConnectionResetError: [Errno 54] Connection reset by peer'
                    in stderr):
                logger.info('Retrying due to Connection reset by peer.')
                return True
            return False

        retry_cnt = 0
        ray_up_return_value = None
        # 5 seconds to 180 seconds. We need backoff for e.g., rate limit per
        # minute errors.
        backoff = common_utils.Backoff(initial_backoff=5,
                                       max_backoff_factor=180 // 5)
        while (retry_cnt < _MAX_RAY_UP_RETRY and
               need_ray_up(ray_up_return_value)):
            retry_cnt += 1
            if retry_cnt > 1:
                sleep = backoff.current_backoff()
                logger.info(
                    'Retrying launching in {:.1f} seconds.'.format(sleep))
                time.sleep(sleep)
            ray_up_return_value = ray_up()

        assert ray_up_return_value is not None
        returncode, stdout, stderr = ray_up_return_value

        logger.debug(f'Ray up takes {time.time() - start} seconds with '
                     f'{retry_cnt} retries.')
        if returncode != 0:
            return self.GangSchedulingStatus.HEAD_FAILED, stdout, stderr, None

        resources = cluster_handle.launched_resources
        if tpu_utils.is_tpu_vm_pod(resources):
            logger.info(f'{style.BRIGHT}Setting up TPU VM Pod workers...'
                        f'{style.RESET_ALL}')
            self._tpu_pod_setup(cluster_config_file, cluster_handle, num_nodes)

        # Only 1 node or head node provisioning failure.
        if num_nodes == 1 and returncode == 0:
            # Optimization: Try parse head ip from 'ray up' stdout.
            # Last line looks like: 'ssh ... <user>@<public head_ip>\n'
            position = stdout.rfind('@')
            # Use a regex to extract the IP address.
            ip_list = re.findall(backend_utils.IP_ADDR_REGEX,
                                 stdout[position + 1:])
            # If something's wrong. Ok to not return a head_ip.
            head_ip = None
            if len(ip_list) == 1:
                head_ip = ip_list[0]
            return (self.GangSchedulingStatus.CLUSTER_READY, stdout, stderr,
                    head_ip)

        # All code below is handling num_nodes > 1.

        provision_str = 'Successfully provisioned or found existing head VM.'
        if isinstance(to_provision_cloud, clouds.Local):
            provision_str = 'Successfully connected to head node.'

        logger.info(f'{style.BRIGHT}{provision_str} '
                    f'Waiting for workers.{style.RESET_ALL}')

        # Special handling is needed for the local case. This is due to a Ray
        # autoscaler bug, where filemounting and setup does not run on worker
        # nodes. Hence, this method here replicates what the Ray autoscaler
        # would do were it for public cloud.
        if isinstance(to_provision_cloud, clouds.Local):
            backend_utils.do_filemounts_and_setup_on_local_workers(
                cluster_config_file)

        # FIXME(zongheng): the below requires ray processes are up on head. To
        # repro it failing: launch a 2-node cluster, log into head and ray
        # stop, then launch again.
        cluster_ready = backend_utils.wait_until_ray_cluster_ready(
            cluster_config_file,
            num_nodes,
            log_path=log_abs_path,
            nodes_launching_progress_timeout=_NODES_LAUNCHING_PROGRESS_TIMEOUT,
            is_local_cloud=isinstance(to_provision_cloud, clouds.Local))
        if cluster_ready:
            cluster_status = self.GangSchedulingStatus.CLUSTER_READY
            # ray up --no-restart again with upscaling_speed=0 after cluster is
            # ready to ensure cluster will not scale up after preemption (spot).
            # Skip for non-spot as this takes extra time to provision (~1min).
            if use_spot:
                ray_config = common_utils.read_yaml(cluster_config_file)
                ray_config['upscaling_speed'] = 0
                common_utils.dump_yaml(cluster_config_file, ray_config)
                start = time.time()
                returncode, stdout, stderr = ray_up()
                logger.debug(
                    f'Upscaling reset takes {time.time() - start} seconds.')
                if returncode != 0:
                    return (self.GangSchedulingStatus.GANG_FAILED, stdout,
                            stderr, None)
        else:
            cluster_status = self.GangSchedulingStatus.GANG_FAILED

        # Do not need stdout/stderr if gang scheduling failed.
        # gang_succeeded = False, if head OK, but workers failed.
        return cluster_status, '', '', None

    def _ensure_cluster_ray_started(self,
                                    handle: 'CloudVmRayBackend.ResourceHandle',
                                    log_abs_path) -> None:
        """Ensures ray processes are up on a just-provisioned cluster."""
        if handle.launched_nodes > 1:
            # FIXME(zongheng): this has NOT been tested with multinode
            # clusters; mainly because this function will not be reached in
            # that case.  See #140 for details.  If it were reached, the
            # following logic might work:
            #   - get all node ips
            #   - for all nodes: ray stop
            #   - ray up --restart-only
            return
        backend = CloudVmRayBackend()

        returncode = backend.run_on_head(
            handle,
            'ray status',
            # At this state, an erroneous cluster may not have cached
            # handle.head_ip (global_user_state.add_or_update_cluster(...,
            # ready=True)).
            use_cached_head_ip=False)
        if returncode == 0:
            return
        launched_resources = handle.launched_resources
        # Ray cluster should already be running if the system admin has setup
        # Ray.
        if isinstance(launched_resources.cloud, clouds.Local):
            raise RuntimeError(
                'The command `ray status` errored out on the head node '
                'of the local cluster. Check if ray[default]==2.0.1 '
                'is installed or running correctly.')
        backend.run_on_head(handle, 'ray stop', use_cached_head_ip=False)

        log_lib.run_with_log(
            ['ray', 'up', '-y', '--restart-only', handle.cluster_yaml],
            log_abs_path,
            stream_logs=False,
            # Use environment variables to disable the ray usage collection
            # (avoid the 10 second wait for usage collection confirmation),
            # as the ray version on the user's machine may be lower version
            # that does not support the `--disable-usage-stats` flag.
            env=dict(os.environ, RAY_USAGE_STATS_ENABLED='0'),
            # Disable stdin to avoid ray outputs mess up the terminal with
            # misaligned output when multithreading/multiprocessing is used.
            # Refer to: https://github.com/ray-project/ray/blob/d462172be7c5779abf37609aed08af112a533e1e/python/ray/autoscaler/_private/subprocess_output_util.py#L264 # pylint: disable=line-too-long
            stdin=subprocess.DEVNULL)

    @timeline.event
    def provision_with_retries(
        self,
        task: task_lib.Task,
        to_provision_config: ToProvisionConfig,
        dryrun: bool,
        stream_logs: bool,
    ):
        """Provision with retries for all launchable resources."""
        cluster_name = to_provision_config.cluster_name
        to_provision = to_provision_config.resources
        num_nodes = to_provision_config.num_nodes
        cluster_exists = to_provision_config.cluster_exists
        launchable_retries_disabled = (self._dag is None or
                                       self._optimize_target is None)

        style = colorama.Style
        # Retrying launchable resources.
        provision_failed = True
        while provision_failed:
            provision_failed = False
            try:
                try:
                    # Recheck cluster name as the 'except:' block below may
                    # change the cloud assignment.
                    backend_utils.check_cluster_name_is_valid(
                        cluster_name, to_provision.cloud)
                except exceptions.InvalidClusterNameError as e:
                    # Let failover below handle this (i.e., block this cloud).
                    raise exceptions.ResourcesUnavailableError(str(e)) from e
                config_dict = self._retry_region_zones(
                    to_provision,
                    num_nodes,
                    dryrun=dryrun,
                    stream_logs=stream_logs,
                    cluster_name=cluster_name,
                    cluster_exists=cluster_exists)
                if dryrun:
                    return
            except exceptions.ResourcesUnavailableError as e:
                if e.no_failover:
                    raise e
                if launchable_retries_disabled:
                    logger.warning(
                        'DAG and optimize_target needs to be registered first '
                        'to enable cross-cloud retry. '
                        'To fix, call backend.register_info(dag=dag, '
                        'optimize_target=sky.OptimizeTarget.COST)')
                    raise e

                logger.warning(e)
                provision_failed = True
                logger.warning(
                    f'\n{style.BRIGHT}Provision failed for {num_nodes}x '
                    f'{to_provision}. Trying other launchable resources '
                    f'(if any).{style.RESET_ALL}')
                if not cluster_exists:
                    # Add failed resources to the blocklist, only when it
                    # is in fallback mode.
                    self._blocked_launchable_resources.add(to_provision)
                else:
                    logger.info(
                        'Retrying provisioning with requested resources '
                        f'{task.num_nodes}x {task.resources}')
                    # Retry with the current, potentially "smaller" resources:
                    # to_provision == the current new resources (e.g., V100:1),
                    # which may be "smaller" than the original (V100:8).
                    # num_nodes is not part of a Resources so must be updated
                    # separately.
                    num_nodes = task.num_nodes
                    cluster_exists = False

                # Set to None so that sky.optimize() will assign a new one
                # (otherwise will skip re-optimizing this task).
                # TODO: set all remaining tasks' best_resources to None.
                task.best_resources = None
                self._dag = sky.optimize(self._dag,
                                         minimize=self._optimize_target,
                                         blocked_launchable_resources=self.
                                         _blocked_launchable_resources)
                to_provision = task.best_resources
                assert task in self._dag.tasks, 'Internal logic error.'
                assert to_provision is not None, task
        return config_dict


class CloudVmRayBackend(backends.Backend):
    """Backend: runs on cloud virtual machines, managed by Ray.

    Changing this class may also require updates to:
      * Cloud providers' templates under config/
      * Cloud providers' implementations under clouds/
    """

    NAME = 'cloudvmray'

    class ResourceHandle(object):
        """A pickle-able tuple of:

        - (required) Cluster name.
        - (required) Path to a cluster.yaml file.
        - (optional) A cached head node public IP.  Filled in after a
            successful provision().
        - (optional) Launched num nodes
        - (optional) Launched resources
        - (optional) If TPU(s) are managed, a path to a deletion script.
        """
        _VERSION = 2

        def __init__(self,
                     *,
                     cluster_name: str,
                     cluster_yaml: str,
                     head_ip: Optional[str] = None,
                     launched_nodes: Optional[int] = None,
                     launched_resources: Optional[
                         resources_lib.Resources] = None,
                     tpu_create_script: Optional[str] = None,
                     tpu_delete_script: Optional[str] = None) -> None:
            self._version = self._VERSION
            self.cluster_name = cluster_name
            self._cluster_yaml = cluster_yaml.replace(os.path.expanduser('~'),
                                                      '~', 1)
            self.head_ip = head_ip
            self.launched_nodes = launched_nodes
            self.launched_resources = launched_resources
            self.tpu_create_script = tpu_create_script
            self.tpu_delete_script = tpu_delete_script
            self._maybe_make_local_handle()

        def __repr__(self):
            return (f'ResourceHandle('
                    f'\n\tcluster_name={self.cluster_name},'
                    f'\n\thead_ip={self.head_ip},'
                    '\n\tcluster_yaml='
                    f'{self.cluster_yaml}, '
                    f'\n\tlaunched_resources={self.launched_nodes}x '
                    f'{self.launched_resources}, '
                    f'\n\ttpu_create_script={self.tpu_create_script}, '
                    f'\n\ttpu_delete_script={self.tpu_delete_script})')

        def get_cluster_name(self):
            return self.cluster_name

        def _maybe_make_local_handle(self):
            """Adds local handle for the local cloud case.

            For public cloud, the first time sky launch is ran, task resources
            = cluster resources. For the local cloud case, sky launch is ran,
            task resources != cluster resources; hence, this method is needed
            to correct this.
            """
            self.local_handle = None
            local_file = os.path.expanduser(
                onprem_utils.SKY_USER_LOCAL_CONFIG_PATH.format(
                    self.cluster_name))
            # Local cluster case requires several modifications:
            #   1) Create local_handle to store local cluster IPs and
            #      custom accelerators for each node.
            #   2) Replace launched_resources to represent a generic local
            #      node (without accelerator specifications).
            #   3) Replace launched_nodes to represent the total nodes in the
            #      local cluster.
            if os.path.isfile(local_file):
                config = onprem_utils.get_local_cluster_config_or_error(
                    self.cluster_name)
                self.local_handle = {}
                cluster_config = config['cluster']
                auth_config = config['auth']
                ips = cluster_config['ips']
                local_region = clouds.Local.regions()[0].name
                # Convert existing ResourceHandle fields to specify local
                # cluster resources.
                self.launched_resources = resources_lib.Resources(
                    cloud=clouds.Local(), region=local_region)
                self.launched_nodes = len(ips)
                self.local_handle['ips'] = ips
                cluster_accs = onprem_utils.get_local_cluster_accelerators(
                    ips, auth_config)
                self.local_handle['cluster_resources'] = [
                    resources_lib.Resources(
                        cloud=clouds.Local(),
                        accelerators=acc_dict if acc_dict else None,
                        region=local_region) for acc_dict in cluster_accs
                ]

        def _update_cluster_region(self):
            if self.launched_resources.region is not None:
                return

            config = common_utils.read_yaml(self.cluster_yaml)
            provider = config['provider']
            cloud = self.launched_resources.cloud
            if cloud.is_same_cloud(clouds.Azure()):
                region = provider['location']
            elif cloud.is_same_cloud(clouds.GCP()) or cloud.is_same_cloud(
                    clouds.AWS()):
                region = provider['region']
            elif cloud.is_same_cloud(clouds.Local()):
                # There is only 1 region for Local cluster, 'Local'.
                local_regions = clouds.Local.regions()
                region = local_regions[0].name

            self.launched_resources = self.launched_resources.copy(
                region=region)

        @property
        def cluster_yaml(self):
            return os.path.expanduser(self._cluster_yaml)

        def __setstate__(self, state):
            self._version = self._VERSION

            version = state.pop('_version', None)
            if version is None:
                version = -1
                state.pop('cluster_region', None)
            if version < 2:
                state['_cluster_yaml'] = state.pop('cluster_yaml')

            self.__dict__.update(state)
            self._update_cluster_region()

    def __init__(self):
        self.run_timestamp = backend_utils.get_run_timestamp()
        self.log_dir = os.path.join(constants.SKY_LOGS_DIRECTORY,
                                    self.run_timestamp)
        # Do not make directories to avoid create folder for commands that
        # do not need it (`sky status`, `sky logs` ...)
        # os.makedirs(self.log_dir, exist_ok=True)

        self._dag = None
        self._optimize_target = None

        # Command for running the setup script. It is only set when the
        # setup needs to be run outside the self._setup() and as part of
        # a job (--detach-setup).
        self._setup_cmd = None

    # --- Implementation of Backend APIs ---

    def register_info(self, **kwargs) -> None:
        self._dag = kwargs.pop('dag', self._dag)
        self._optimize_target = kwargs.pop(
            'optimize_target', self._optimize_target) or OptimizeTarget.COST
        assert len(kwargs) == 0, f'Unexpected kwargs: {kwargs}'

    def check_resources_fit_cluster(self, handle: ResourceHandle,
                                    task: task_lib.Task):
        """Check if resources requested by the task fit the cluster.

        The resources requested by the task should be smaller than the existing
        cluster.
        """
        assert len(task.resources) == 1, task.resources

        launched_resources = handle.launched_resources
        task_resources = list(task.resources)[0]
        cluster_name = handle.cluster_name
        usage_lib.messages.usage.update_cluster_resources(
            handle.launched_nodes, launched_resources)
        record = global_user_state.get_cluster_from_name(cluster_name)
        if record is not None:
            usage_lib.messages.usage.update_cluster_status(record['status'])

        # Backward compatibility: the old launched_resources without region info
        # was handled by ResourceHandle._update_cluster_region.
        assert launched_resources.region is not None, handle

        mismatch_str = (f'To fix: specify a new cluster name, or down the '
                        f'existing cluster first: sky down {cluster_name}')
        if hasattr(handle, 'local_handle') and handle.local_handle is not None:
            launched_resources = handle.local_handle['cluster_resources']
            usage_lib.messages.usage.update_local_cluster_resources(
                launched_resources)
            mismatch_str = ('To fix: use accelerators/number of nodes that can '
                            'be satisfied by the local cluster')
        # Requested_resources <= actual_resources.
        # Special handling for local cloud case, which assumes a cluster can
        # be heterogeneous. Here, launched_resources is a list of custom
        # accelerators per node, and Resources.less_demanding_than determines
        # how many nodes satisfy task resource requirements.
        if not (task.num_nodes <= handle.launched_nodes and
                task_resources.less_demanding_than(
                    launched_resources, requested_num_nodes=task.num_nodes)):
            if (task_resources.region is not None and
                    task_resources.region != launched_resources.region):
                with ux_utils.print_exception_no_traceback():
                    raise exceptions.ResourcesMismatchError(
                        'Task requested resources in region '
                        f'{task_resources.region!r}, but the existing cluster '
                        f'is in region {launched_resources.region!r}.')
            if (task_resources.zone is not None and
                    task_resources.zone != launched_resources.zone):
                zone_str = (f'is in zone {launched_resources.zone!r}.'
                            if launched_resources.zone is not None else
                            'does not have zone specified.')
                with ux_utils.print_exception_no_traceback():
                    raise exceptions.ResourcesMismatchError(
                        'Task requested resources in zone '
                        f'{task_resources.zone!r}, but the existing cluster '
                        f'{zone_str}')
            with ux_utils.print_exception_no_traceback():
                raise exceptions.ResourcesMismatchError(
                    'Requested resources do not match the existing cluster.\n'
                    f'  Requested:\t{task.num_nodes}x {task_resources} \n'
                    f'  Existing:\t{handle.launched_nodes}x '
                    f'{handle.launched_resources}\n'
                    f'{mismatch_str}')

    def _provision(self,
                   task: task_lib.Task,
                   to_provision: Optional[resources_lib.Resources],
                   dryrun: bool,
                   stream_logs: bool,
                   cluster_name: str,
                   retry_until_up: bool = False) -> ResourceHandle:
        """Provisions using 'ray up'."""
        # FIXME: ray up for Azure with different cluster_names will overwrite
        # each other.
        lock_path = os.path.expanduser(
            backend_utils.CLUSTER_STATUS_LOCK_PATH.format(cluster_name))
        with timeline.FileLockEvent(lock_path):
            to_provision_config = RetryingVmProvisioner.ToProvisionConfig(
                cluster_name, to_provision, task.num_nodes)
            prev_cluster_status = None
            if not dryrun:  # dry run doesn't need to check existing cluster.
                # Try to launch the exiting cluster first
                to_provision_config = self._check_existing_cluster(
                    task, to_provision, cluster_name)
                prev_cluster_status, _ = (
                    backend_utils.refresh_cluster_status_handle(
                        cluster_name, acquire_per_cluster_status_lock=False))
            assert to_provision_config.resources is not None, (
                'to_provision should not be None', to_provision_config)

            usage_lib.messages.usage.update_cluster_resources(
                to_provision_config.num_nodes, to_provision_config.resources)
            usage_lib.messages.usage.update_cluster_status(prev_cluster_status)

            # TODO(suquark): once we have sky on PyPI, we should directly
            # install sky from PyPI.
            # NOTE: can take ~2s.
            with timeline.Event('backend.provision.wheel_build'):
                # TODO(suquark): once we have sky on PyPI, we should directly
                # install sky from PyPI.
                local_wheel_path, wheel_hash = wheel_utils.build_sky_wheel()
            backoff = common_utils.Backoff(_RETRY_UNTIL_UP_INIT_GAP_SECONDS)
            attempt_cnt = 1
            while True:
                # RetryingVmProvisioner will retry within a cloud's regions
                # first (if a region is not explicitly requested), then
                # optionally retry on all other clouds (if
                # backend.register_info() has been called).  After this "round"
                # of optimization across clouds, provisioning may still have
                # not succeeded. This while loop will then kick in if
                # retry_until_up is set, which will kick off new "rounds" of
                # optimization infinitely.
                try:
                    provisioner = RetryingVmProvisioner(self.log_dir, self._dag,
                                                        self._optimize_target,
                                                        local_wheel_path,
                                                        wheel_hash)
                    config_dict = provisioner.provision_with_retries(
                        task, to_provision_config, dryrun, stream_logs)
                    break
                except exceptions.ResourcesUnavailableError as e:
                    # Do not remove the stopped cluster from the global state
                    # if failed to start.
                    if e.no_failover:
                        error_message = str(e)
                    else:
                        # Clean up the cluster's entry in `sky status`.
                        global_user_state.remove_cluster(cluster_name,
                                                         terminate=True)
                        usage_lib.messages.usage.update_final_cluster_status(
                            None)
                        error_message = (
                            'Failed to provision all possible launchable '
                            'resources.'
                            f' Relax the task\'s resource requirements: '
                            f'{task.num_nodes}x {task.resources}')
                    if retry_until_up:
                        logger.error(error_message)
                        # Sleep and retry.
                        gap_seconds = backoff.current_backoff()
                        plural = 's' if attempt_cnt > 1 else ''
                        logger.info(
                            f'{colorama.Style.BRIGHT}=== Retry until up ==='
                            f'{colorama.Style.RESET_ALL}\n'
                            f'Retrying provisioning after {gap_seconds:.0f}s '
                            '(exponential backoff with random jittering). '
                            f'Already tried {attempt_cnt} attempt{plural}.')
                        attempt_cnt += 1
                        time.sleep(gap_seconds)
                        continue
                    error_message += (
                        '\nTo keep retrying until the cluster is up, use the '
                        '`--retry-until-up` flag.')
                    with ux_utils.print_exception_no_traceback():
                        raise exceptions.ResourcesUnavailableError(
                            error_message) from None
            if dryrun:
                return
            cluster_config_file = config_dict['ray']

            if 'tpu_name' in config_dict:
                self._set_tpu_name(cluster_config_file,
                                   config_dict['launched_nodes'],
                                   config_dict['tpu_name'])

            if config_dict['launched_nodes'] == 1 and config_dict[
                    'head_ip'] is not None:
                # Optimization for 1-node: we may have parsed the stdout of
                # 'ray up' to get the head_ip already.
                head_ip = config_dict['head_ip']
                ip_list = [head_ip]
            else:
                # NOTE: querying node_ips is expensive, observed 1node GCP >=4s.
                with timeline.Event('backend.provision.get_node_ips'):
                    ip_list = backend_utils.get_node_ips(
                        cluster_config_file,
                        config_dict['launched_nodes'],
                        head_ip_max_attempts=_HEAD_IP_MAX_ATTEMPTS,
                        worker_ip_max_attempts=_WORKER_IP_MAX_ATTEMPTS)
                    head_ip = ip_list[0]

            handle = self.ResourceHandle(
                cluster_name=cluster_name,
                cluster_yaml=cluster_config_file,
                # Cache head ip in the handle to speed up ssh operations.
                head_ip=head_ip,
                launched_nodes=config_dict['launched_nodes'],
                launched_resources=config_dict['launched_resources'],
                # TPU.
                tpu_create_script=config_dict.get('tpu-create-script'),
                tpu_delete_script=config_dict.get('tpu-delete-script'))

            # Get actual zone info and save it into handle.
            # NOTE: querying zones is expensive, observed 1node GCP >=4s.
            zones = config_dict['zones']
            if zones is not None and len(zones) == 1:  # zones is None for Azure
                # Optimization for if the provision request was for 1 zone
                # (currently happens only for GCP since it uses per-zone
                # provisioning), then we know the exact zone already.
                handle.launched_resources = handle.launched_resources.copy(
                    zone=zones[0].name)
            elif (task.num_nodes == 1 or
                  handle.launched_resources.zone is not None):
                # Query zone if the cluster has 1 node, or a zone is
                # specifically requested for a multinode cluster.  Otherwise
                # leave the zone field to None because head and worker nodes
                # can be launched in different zones.
                get_zone_cmd = (
                    handle.launched_resources.cloud.get_zone_shell_cmd())
                if get_zone_cmd is not None:
                    returncode, stdout, stderr = self.run_on_head(
                        handle, get_zone_cmd, require_outputs=True)
                    subprocess_utils.handle_returncode(returncode,
                                                       get_zone_cmd,
                                                       'Failed to get zone',
                                                       stderr=stderr,
                                                       stream_logs=stream_logs)
                    # zone will be checked during Resources cls
                    # initialization.
                    handle.launched_resources = handle.launched_resources.copy(
                        zone=stdout.strip())

            usage_lib.messages.usage.update_cluster_resources(
                handle.launched_nodes, handle.launched_resources)
            usage_lib.messages.usage.update_final_cluster_status(
                global_user_state.ClusterStatus.UP)

            # Update job queue to avoid stale jobs (when restarted), before
            # setting the cluster to be ready.
            if prev_cluster_status == global_user_state.ClusterStatus.INIT:
                # update_status will query the ray job status for all INIT /
                # PENDING / RUNNING jobs for the real status, since we do not
                # know the actual previous status of the cluster.
                job_owner = onprem_utils.get_job_owner(handle.cluster_yaml)
                cmd = job_lib.JobLibCodeGen.update_status(job_owner)
                with backend_utils.safe_console_status(
                        '[bold cyan]Preparing Job Queue'):
                    returncode, _, stderr = self.run_on_head(
                        handle, cmd, require_outputs=True)
                subprocess_utils.handle_returncode(
                    returncode, cmd, 'Failed to update job status.', stderr)
            if prev_cluster_status == global_user_state.ClusterStatus.STOPPED:
                # Safely set all the previous jobs to FAILED since the cluster
                # is restarted
                # An edge case here due to racing:
                # 1. A job finishes RUNNING, but right before it update itself
                # to SUCCEEDED, the cluster is STOPPED by `sky stop`.
                # 2. On next `sky start`, it gets reset to FAILED.
                cmd = job_lib.JobLibCodeGen.fail_all_jobs_in_progress()
                returncode, stdout, stderr = self.run_on_head(
                    handle, cmd, require_outputs=True)
                subprocess_utils.handle_returncode(
                    returncode, cmd,
                    'Failed to set previously in-progress jobs to FAILED',
                    stdout + stderr)

            with timeline.Event('backend.provision.post_process'):
                global_user_state.add_or_update_cluster(cluster_name,
                                                        handle,
                                                        ready=True)
                usage_lib.messages.usage.update_final_cluster_status(
                    global_user_state.ClusterStatus.UP)
                auth_config = common_utils.read_yaml(
                    handle.cluster_yaml)['auth']
                backend_utils.SSHConfigHelper.add_cluster(
                    cluster_name, ip_list, auth_config)

                os.remove(lock_path)
                return handle

    def _sync_workdir(self, handle: ResourceHandle, workdir: Path) -> None:
        # Even though provision() takes care of it, there may be cases where
        # this function is called in isolation, without calling provision(),
        # e.g., in CLI.  So we should rerun rsync_up.
        fore = colorama.Fore
        style = colorama.Style
        ip_list = backend_utils.get_node_ips(handle.cluster_yaml,
                                             handle.launched_nodes,
                                             handle=handle)
        full_workdir = os.path.abspath(os.path.expanduser(workdir))

        # These asserts have been validated at Task construction time.
        assert os.path.exists(full_workdir), f'{full_workdir} does not exist'
        if os.path.islink(full_workdir):
            logger.warning(
                f'{fore.YELLOW}Workdir {workdir!r} is a symlink. '
                f'Symlink contents are not uploaded.{style.RESET_ALL}')
        else:
            assert os.path.isdir(
                full_workdir), f'{full_workdir} should be a directory.'

        # Raise warning if directory is too large
        dir_size = backend_utils.path_size_megabytes(full_workdir)
        if dir_size >= _PATH_SIZE_MEGABYTES_WARN_THRESHOLD:
            logger.warning(
                f'{fore.YELLOW}The size of workdir {workdir!r} '
                f'is {dir_size} MB. Try to keep workdir small or use '
                '.gitignore to exclude large files, as large sizes will slow '
                f'down rsync.{style.RESET_ALL}')

        log_path = os.path.join(self.log_dir, 'workdir_sync.log')

        ssh_credentials = backend_utils.ssh_credential_from_yaml(
            handle.cluster_yaml)

        # TODO(zhwu): refactor this with backend_utils.parallel_cmd_with_rsync
        runners = command_runner.SSHCommandRunner.make_runner_list(
            ip_list, **ssh_credentials)

        def _sync_workdir_node(runner: command_runner.SSHCommandRunner) -> None:
            runner.rsync(
                source=workdir,
                target=SKY_REMOTE_WORKDIR,
                up=True,
                log_path=log_path,
                stream_logs=False,
            )

        num_nodes = handle.launched_nodes
        plural = 's' if num_nodes > 1 else ''
        logger.info(
            f'{fore.CYAN}Syncing workdir (to {num_nodes} node{plural}): '
            f'{style.BRIGHT}{workdir}{style.RESET_ALL}'
            f' -> '
            f'{style.BRIGHT}{SKY_REMOTE_WORKDIR}{style.RESET_ALL}')
        tail_cmd = f'tail -n100 -f {log_path}'
        logger.info('To view detailed progress: '
                    f'{style.BRIGHT}{tail_cmd}{style.RESET_ALL}')
        with backend_utils.safe_console_status('[bold cyan]Syncing[/]'):
            subprocess_utils.run_in_parallel(_sync_workdir_node, runners)

    def _sync_file_mounts(
        self,
        handle: ResourceHandle,
        all_file_mounts: Dict[Path, Path],
        storage_mounts: Dict[Path, storage_lib.Storage],
    ) -> None:
        """Mounts all user files to the remote nodes."""
        self._execute_file_mounts(handle, all_file_mounts)
        self._execute_storage_mounts(handle, storage_mounts)

    def _setup(self, handle: ResourceHandle, task: task_lib.Task,
               detach_setup: bool) -> Optional[str]:
        start = time.time()
        style = colorama.Style
        fore = colorama.Fore

        if task.setup is None:
            return

        setup_script = log_lib.make_task_bash_script(task.setup,
                                                     env_vars=task.envs)
        with tempfile.NamedTemporaryFile('w', prefix='sky_setup_') as f:
            f.write(setup_script)
            f.flush()
            setup_sh_path = f.name
            setup_file = os.path.basename(setup_sh_path)
            # Sync the setup script up and run it.
            ip_list = backend_utils.get_node_ips(
                handle.cluster_yaml,
                handle.launched_nodes,
                handle=handle,
                head_ip_max_attempts=_HEAD_IP_MAX_ATTEMPTS,
                worker_ip_max_attempts=_WORKER_IP_MAX_ATTEMPTS)
            ssh_credentials = backend_utils.ssh_credential_from_yaml(
                handle.cluster_yaml)
            # Disable connection sharing for setup script to avoid old
            # connections being reused, which may cause stale ssh agent
            # forwarding.
            ssh_credentials.pop('ssh_control_name')
            runners = command_runner.SSHCommandRunner.make_runner_list(
                ip_list, **ssh_credentials)

            # Need this `-i` option to make sure `source ~/.bashrc` work
            setup_cmd = f'/bin/bash -i /tmp/{setup_file} 2>&1'

            def _setup_node(runner: command_runner.SSHCommandRunner) -> int:
                runner.rsync(source=setup_sh_path,
                             target=f'/tmp/{setup_file}',
                             up=True,
                             stream_logs=False)
                if detach_setup:
                    return
                setup_log_path = os.path.join(self.log_dir,
                                              f'setup-{runner.ip}.log')
                returncode = runner.run(
                    setup_cmd,
                    log_path=setup_log_path,
                    process_stream=False,
                )

                def error_message() -> str:
                    # Use the function to avoid tailing the file in success case
                    try:
                        last_10_lines = subprocess.run(
                            [
                                'tail', '-n10',
                                os.path.expanduser(setup_log_path)
                            ],
                            stdout=subprocess.PIPE,
                            check=True).stdout.decode('utf-8')
                    except subprocess.CalledProcessError:
                        last_10_lines = None

                    err_msg = (
                        f'Failed to setup with return code {returncode}. '
                        f'Check the details in log: {setup_log_path}')
                    if last_10_lines:
                        err_msg += (
                            f'\n\n{colorama.Fore.RED}'
                            '****** START Last lines of setup output ******'
                            f'{colorama.Style.RESET_ALL}\n'
                            f'{last_10_lines}'
                            f'{colorama.Fore.RED}'
                            '******* END Last lines of setup output *******'
                            f'{colorama.Style.RESET_ALL}')
                    return err_msg

                subprocess_utils.handle_returncode(returncode=returncode,
                                                   command=setup_cmd,
                                                   error_msg=error_message)

            num_nodes = len(ip_list)
            plural = 's' if num_nodes > 1 else ''
            if not detach_setup:
                logger.info(
                    f'{fore.CYAN}Running setup on {num_nodes} node{plural}.'
                    f'{style.RESET_ALL}')
            subprocess_utils.run_in_parallel(_setup_node, runners)

        if detach_setup:
            # Only set this when setup needs to be run outside the self._setup()
            # as part of a job (--detach-setup).
            self._setup_cmd = setup_cmd
            return
        logger.info(f'{fore.GREEN}Setup completed.{style.RESET_ALL}')
        end = time.time()
        logger.debug(f'Setup took {end - start} seconds.')

    def _exec_code_on_head(
        self,
        handle: ResourceHandle,
        codegen: str,
        job_id: int,
        executable: str,
        detach_run: bool = False,
    ) -> None:
        """Executes generated code on the head node."""
        colorama.init()
        style = colorama.Style
        fore = colorama.Fore
        ssh_credentials = backend_utils.ssh_credential_from_yaml(
            handle.cluster_yaml)
        runner = command_runner.SSHCommandRunner(handle.head_ip,
                                                 **ssh_credentials)
        with tempfile.NamedTemporaryFile('w', prefix='sky_app_') as fp:
            fp.write(codegen)
            fp.flush()
            script_path = os.path.join(SKY_REMOTE_APP_DIR, f'sky_job_{job_id}')
            # We choose to sync code + exec, because the alternative of 'ray
            # submit' may not work as it may use system python (python2) to
            # execute the script.  Happens for AWS.
            runner.rsync(source=fp.name,
                         target=script_path,
                         up=True,
                         stream_logs=False)
        remote_log_dir = self.log_dir
        remote_log_path = os.path.join(remote_log_dir, 'run.log')

        assert executable == 'python3', executable
        cd = f'cd {SKY_REMOTE_WORKDIR}'

        ray_job_id = job_lib.make_ray_job_id(job_id,
                                             ssh_credentials['ssh_user'])
        if isinstance(handle.launched_resources.cloud, clouds.Local):
            # Ray Multitenancy is unsupported.
            # (Git Issue) https://github.com/ray-project/ray/issues/6800
            # Temporary workaround - wrap the run command in a script, and
            # execute it as the specified user.
            executable = onprem_utils.get_python_executable(handle.cluster_name)
            ray_command = (f'{cd} && {executable} -u {script_path} '
                           f'> {remote_log_path} 2>&1')
            job_submit_cmd = self._setup_and_create_job_cmd_on_local_head(
                handle, ray_command, ray_job_id)
        else:
            job_submit_cmd = (
                f'{cd} && mkdir -p {remote_log_dir} && ray job submit '
                f'--address=http://127.0.0.1:8265 --submission-id {ray_job_id} '
                '--no-wait '
                f'"{executable} -u {script_path} > {remote_log_path} 2>&1"')

        returncode, stdout, stderr = self.run_on_head(handle,
                                                      job_submit_cmd,
                                                      stream_logs=False,
                                                      require_outputs=True)
        subprocess_utils.handle_returncode(returncode,
                                           job_submit_cmd,
                                           f'Failed to submit job {job_id}.',
                                           stderr=stdout + stderr)

        logger.info('Job submitted with Job ID: '
                    f'{style.BRIGHT}{job_id}{style.RESET_ALL}')

        try:
            if not detach_run:
                if handle.cluster_name == spot_lib.SPOT_CONTROLLER_NAME:
                    self.tail_spot_logs(handle, job_id)
                else:
                    # Sky logs. Not using subprocess.run since it will make the
                    # ssh keep connected after ctrl-c.
                    self.tail_logs(handle, job_id)
        finally:
            name = handle.cluster_name
            if name == spot_lib.SPOT_CONTROLLER_NAME:
                logger.info(f'{fore.CYAN}Spot Job ID: '
                            f'{style.BRIGHT}{job_id}{style.RESET_ALL}'
                            '\nTo cancel the job:\t\t'
                            f'{backend_utils.BOLD}sky spot cancel {job_id}'
                            f'{backend_utils.RESET_BOLD}'
                            '\nTo stream the logs:\t\t'
                            f'{backend_utils.BOLD}sky spot logs {job_id}'
                            f'{backend_utils.RESET_BOLD}'
                            f'\nTo stream controller logs:\t'
                            f'{backend_utils.BOLD}sky logs {name} {job_id}'
                            f'{backend_utils.RESET_BOLD}'
                            '\nTo view all spot jobs:\t\t'
                            f'{backend_utils.BOLD}sky spot queue'
                            f'{backend_utils.RESET_BOLD}')
            else:
                logger.info(f'{fore.CYAN}Job ID: '
                            f'{style.BRIGHT}{job_id}{style.RESET_ALL}'
                            '\nTo cancel the job:\t'
                            f'{backend_utils.BOLD}sky cancel {name} {job_id}'
                            f'{backend_utils.RESET_BOLD}'
                            '\nTo stream the logs:\t'
                            f'{backend_utils.BOLD}sky logs {name} {job_id}'
                            f'{backend_utils.RESET_BOLD}'
                            '\nTo view the job queue:\t'
                            f'{backend_utils.BOLD}sky queue {name}'
                            f'{backend_utils.RESET_BOLD}')

    def _setup_and_create_job_cmd_on_local_head(
        self,
        handle: ResourceHandle,
        ray_command: str,
        ray_job_id: str,
    ):
        """Generates and prepares job submission code for local clusters."""
        ssh_credentials = backend_utils.ssh_credential_from_yaml(
            handle.cluster_yaml)
        ssh_user = ssh_credentials['ssh_user']
        runner = command_runner.SSHCommandRunner(handle.head_ip,
                                                 **ssh_credentials)
        remote_log_dir = self.log_dir
        with tempfile.NamedTemporaryFile('w', prefix='sky_local_app_') as fp:
            fp.write(ray_command)
            fp.flush()
            run_file = os.path.basename(fp.name)
            remote_run_file = f'/tmp/{run_file}'
            # We choose to sync code + exec, so that Ray job submission API will
            # work for the multitenant case.
            runner.rsync(source=fp.name,
                         target=remote_run_file,
                         up=True,
                         stream_logs=False)
        runner.run(f'mkdir -p {remote_log_dir}; chmod a+rwx {remote_run_file}',
                   stream_logs=False)
        switch_user_cmd = job_lib.make_job_command_with_user_switching(
            ssh_user, remote_run_file)
        switch_user_cmd = ' '.join(switch_user_cmd)
        job_submit_cmd = (
            'ray job submit '
            f'--address=http://127.0.0.1:8265 --submission-id {ray_job_id} '
            f'--no-wait -- {switch_user_cmd}')
        return job_submit_cmd

    def _add_job(self, handle: ResourceHandle, job_name: str,
                 resources_str: str) -> int:
        username = getpass.getuser()
        code = job_lib.JobLibCodeGen.add_job(job_name, username,
                                             self.run_timestamp, resources_str)
        returncode, job_id_str, stderr = self.run_on_head(handle,
                                                          code,
                                                          stream_logs=False,
                                                          require_outputs=True,
                                                          separate_stderr=True)
        # TODO(zhwu): this sometimes will unexpectedly fail, we can add
        # retry for this, after we figure out the reason.
        subprocess_utils.handle_returncode(returncode, code,
                                           'Failed to fetch job id.', stderr)
        try:
            job_id_match = _JOB_ID_PATTERN.search(job_id_str)
            if job_id_match is not None:
                job_id = int(job_id_match.group(1))
            else:
                # For backward compatibility.
                job_id = int(job_id_str)
        except ValueError as e:
            logger.error(stderr)
            raise ValueError(f'Failed to parse job id: {job_id_str}; '
                             f'Returncode: {returncode}') from e
        return job_id

    def _execute(
        self,
        handle: ResourceHandle,
        task: task_lib.Task,
        detach_run: bool,
    ) -> None:
        if task.run is None:
            logger.info('Run commands not specified or empty.')
            return
        # Check the task resources vs the cluster resources. Since `sky exec`
        # will not run the provision and _check_existing_cluster
        self.check_resources_fit_cluster(handle, task)

        resources_str = backend_utils.get_task_resources_str(task)
        job_id = self._add_job(handle, task.name, resources_str)

        is_tpu_vm_pod = tpu_utils.is_tpu_vm_pod(handle.launched_resources)
        # Case: task_lib.Task(run, num_nodes=N) or TPU VM Pods
        if task.num_nodes > 1 or is_tpu_vm_pod:
            self._execute_task_n_nodes(handle, task, job_id, detach_run)
        else:
            # Case: task_lib.Task(run, num_nodes=1)
            self._execute_task_one_node(handle, task, job_id, detach_run)

    def _post_execute(self, handle: ResourceHandle, down: bool) -> None:
        colorama.init()
        fore = colorama.Fore
        style = colorama.Style
        name = handle.cluster_name
        if name == spot_lib.SPOT_CONTROLLER_NAME or down:
            return
        stop_str = ('\nTo stop the cluster:'
                    f'\t{backend_utils.BOLD}sky stop {name}'
                    f'{backend_utils.RESET_BOLD}')
        if isinstance(handle.launched_resources.cloud, clouds.Local):
            stop_str = ''
        logger.info(f'\n{fore.CYAN}Cluster name: '
                    f'{style.BRIGHT}{name}{style.RESET_ALL}'
                    '\nTo log into the head VM:\t'
                    f'{backend_utils.BOLD}ssh {name}'
                    f'{backend_utils.RESET_BOLD}'
                    '\nTo submit a job:'
                    f'\t\t{backend_utils.BOLD}sky exec {name} yaml_file'
                    f'{backend_utils.RESET_BOLD}'
                    f'{stop_str}'
                    '\nTo teardown the cluster:'
                    f'\t{backend_utils.BOLD}sky down {name}'
                    f'{backend_utils.RESET_BOLD}')
        if handle.tpu_delete_script is not None:
            logger.info('Tip: `sky down` will delete launched TPU(s) too.')

    def _teardown_ephemeral_storage(self, task: task_lib.Task) -> None:
        storage_mounts = task.storage_mounts
        if storage_mounts is not None:
            for _, storage in storage_mounts.items():
                if not storage.persistent:
                    storage.delete()

    def _teardown(self,
                  handle: ResourceHandle,
                  terminate: bool,
                  purge: bool = False):
        cluster_name = handle.cluster_name
        lock_path = os.path.expanduser(
            backend_utils.CLUSTER_STATUS_LOCK_PATH.format(cluster_name))

        try:
            # TODO(mraheja): remove pylint disabling when filelock
            # version updated
            # pylint: disable=abstract-class-instantiated
            with filelock.FileLock(
                    lock_path,
                    backend_utils.CLUSTER_STATUS_LOCK_TIMEOUT_SECONDS):
                success = self.teardown_no_lock(handle, terminate, purge)
            if success and terminate:
                os.remove(lock_path)
        except filelock.Timeout as e:
            raise RuntimeError(
                f'Cluster {cluster_name!r} is locked by {lock_path}. '
                'Check to see if it is still being launched.') from e

    # --- CloudVMRayBackend Specific APIs ---

    def get_job_status(
        self,
        handle: ResourceHandle,
        job_ids: Optional[List[str]] = None,
        stream_logs: bool = True
    ) -> Dict[Optional[str], Optional[job_lib.JobStatus]]:
        code = job_lib.JobLibCodeGen.get_job_status(job_ids)
        returncode, stdout, stderr = self.run_on_head(handle,
                                                      code,
                                                      stream_logs=stream_logs,
                                                      require_outputs=True,
                                                      separate_stderr=True)
        subprocess_utils.handle_returncode(returncode, code,
                                           'Failed to get job status.', stderr)
        statuses = job_lib.load_statuses_payload(stdout)
        return statuses

    def cancel_jobs(self, handle: ResourceHandle, jobs: Optional[List[int]]):
        job_owner = onprem_utils.get_job_owner(handle.cluster_yaml)
        code = job_lib.JobLibCodeGen.cancel_jobs(job_owner, jobs)

        # All error messages should have been redirected to stdout.
        returncode, stdout, _ = self.run_on_head(handle,
                                                 code,
                                                 stream_logs=False,
                                                 require_outputs=True)
        subprocess_utils.handle_returncode(
            returncode, code,
            f'Failed to cancel jobs on cluster {handle.cluster_name}.', stdout)

    def sync_down_logs(
            self,
            handle: ResourceHandle,
            job_ids: Optional[str],
            local_dir: str = constants.SKY_LOGS_DIRECTORY) -> Dict[str, str]:
        """Sync down logs for the given job_ids.

        Returns:
            A dictionary mapping job_id to log path.
        """
        code = job_lib.JobLibCodeGen.get_run_timestamp_with_globbing(job_ids)
        returncode, run_timestamps, stderr = self.run_on_head(
            handle,
            code,
            stream_logs=False,
            require_outputs=True,
            separate_stderr=True)
        subprocess_utils.handle_returncode(returncode, code,
                                           'Failed to sync logs.', stderr)
        run_timestamps = common_utils.decode_payload(run_timestamps)
        if not run_timestamps:
            logger.info(f'{colorama.Fore.YELLOW}'
                        'No matching log directories found'
                        f'{colorama.Style.RESET_ALL}')
            return

        job_ids = list(run_timestamps.keys())
        run_timestamps = list(run_timestamps.values())
        remote_log_dirs = [
            os.path.join(constants.SKY_LOGS_DIRECTORY, run_timestamp)
            for run_timestamp in run_timestamps
        ]
        local_log_dirs = [
            os.path.expanduser(os.path.join(local_dir, run_timestamp))
            for run_timestamp in run_timestamps
        ]

        style = colorama.Style
        fore = colorama.Fore
        for job_id, log_dir in zip(job_ids, local_log_dirs):
            logger.info(f'{fore.CYAN}Job {job_id} logs: {log_dir}'
                        f'{style.RESET_ALL}')

        ip_list = backend_utils.get_node_ips(handle.cluster_yaml,
                                             handle.launched_nodes,
                                             handle=handle)
        ssh_credentials = backend_utils.ssh_credential_from_yaml(
            handle.cluster_yaml)
        runners = command_runner.SSHCommandRunner.make_runner_list(
            ip_list, **ssh_credentials)

        def _rsync_down(args) -> None:
            """Rsync down logs from remote nodes.

            Args:
                args: A tuple of (runner, local_log_dir, remote_log_dir)
            """
            (runner, local_log_dir, remote_log_dir) = args
            try:
                os.makedirs(local_log_dir, exist_ok=True)
                runner.rsync(
                    source=f'{remote_log_dir}/*',
                    target=local_log_dir,
                    up=False,
                    stream_logs=False,
                )
            except exceptions.CommandError as e:
                if e.returncode == exceptions.RSYNC_FILE_NOT_FOUND_CODE:
                    # Raised by rsync_down. Remote log dir may not exist, since
                    # the job can be run on some part of the nodes.
                    logger.debug(f'{runner.ip} does not have the tasks/*.')
                else:
                    raise

        parallel_args = [[runner, *item]
                         for item in zip(local_log_dirs, remote_log_dirs)
                         for runner in runners]
        subprocess_utils.run_in_parallel(_rsync_down, parallel_args)
        return dict(zip(job_ids, local_log_dirs))

    def tail_logs(self,
                  handle: ResourceHandle,
                  job_id: Optional[int],
                  spot_job_id: Optional[int] = None,
                  follow: bool = True) -> int:
        job_owner = onprem_utils.get_job_owner(handle.cluster_yaml)
        code = job_lib.JobLibCodeGen.tail_logs(job_owner,
                                               job_id,
                                               spot_job_id=spot_job_id,
                                               follow=follow)
        if job_id is None and spot_job_id is None:
            logger.info(
                'Job ID not provided. Streaming the logs of the latest job.')

        # With the stdin=subprocess.DEVNULL, the ctrl-c will not directly
        # kill the process, so we need to handle it manually here.
        signal.signal(signal.SIGINT, backend_utils.interrupt_handler)
        signal.signal(signal.SIGTSTP, backend_utils.stop_handler)
        try:
            returncode = self.run_on_head(
                handle,
                code,
                stream_logs=True,
                process_stream=False,
                # Allocate a pseudo-terminal to disable output buffering.
                # Otherwise, there may be 5 minutes delay in logging.
                ssh_mode=command_runner.SshMode.INTERACTIVE,
                # Disable stdin to avoid ray outputs mess up the terminal with
                # misaligned output in multithreading/multiprocessing.
                # Refer to: https://github.com/ray-project/ray/blob/d462172be7c5779abf37609aed08af112a533e1e/python/ray/autoscaler/_private/subprocess_output_util.py#L264 # pylint: disable=line-too-long
                stdin=subprocess.DEVNULL,
            )
        except SystemExit as e:
            returncode = e.code
        return returncode

    def tail_spot_logs(self,
                       handle: ResourceHandle,
                       job_id: Optional[int] = None,
                       job_name: Optional[str] = None,
                       follow: bool = True) -> None:
        # if job_name is not None, job_id should be None
        assert job_name is None or job_id is None, (job_name, job_id)
        if job_name is not None:
            code = spot_lib.SpotCodeGen.stream_logs_by_name(job_name, follow)
        else:
            code = spot_lib.SpotCodeGen.stream_logs_by_id(job_id, follow)

        # With the stdin=subprocess.DEVNULL, the ctrl-c will not directly
        # kill the process, so we need to handle it manually here.
        signal.signal(signal.SIGINT, backend_utils.interrupt_handler)
        signal.signal(signal.SIGTSTP, backend_utils.stop_handler)

        # Refer to the notes in tail_logs.
        self.run_on_head(
            handle,
            code,
            stream_logs=True,
            process_stream=False,
            ssh_mode=command_runner.SshMode.INTERACTIVE,
            stdin=subprocess.DEVNULL,
        )

    def teardown_no_lock(self,
                         handle: ResourceHandle,
                         terminate: bool,
                         purge: bool = False,
                         post_teardown_cleanup: bool = True) -> bool:
        """Teardown the cluster without acquiring the cluster status lock.

        NOTE: This method should not be called without holding the cluster
        status lock already.
        """
        log_path = os.path.join(os.path.expanduser(self.log_dir),
                                'teardown.log')
        log_abs_path = os.path.abspath(log_path)
        cloud = handle.launched_resources.cloud
        config = common_utils.read_yaml(handle.cluster_yaml)
        prev_status, _ = backend_utils.refresh_cluster_status_handle(
            handle.cluster_name, acquire_per_cluster_status_lock=False)
        cluster_name = handle.cluster_name
        use_tpu_vm = config['provider'].get('_has_tpus', False)
        if terminate and isinstance(cloud, clouds.Azure):
            # Here we handle termination of Azure by ourselves instead of Ray
            # autoscaler.
            resource_group = config['provider']['resource_group']
            terminate_cmd = f'az group delete -y --name {resource_group}'
            with backend_utils.safe_console_status(f'[bold cyan]Terminating '
                                                   f'[green]{cluster_name}'):
                returncode, stdout, stderr = log_lib.run_with_log(
                    terminate_cmd,
                    log_abs_path,
                    shell=True,
                    stream_logs=False,
                    require_outputs=True)
        elif (terminate and
              (prev_status == global_user_state.ClusterStatus.STOPPED or
               use_tpu_vm)):
            # For TPU VMs, gcloud CLI is used for VM termination.
            if isinstance(cloud, clouds.AWS):
                # TODO(zhwu): Room for optimization. We can move these cloud
                # specific handling to the cloud class.
                # The stopped instance on AWS will not be correctly terminated
                # due to ray's bug.
                region = config['provider']['region']
                query_cmd = (
                    f'aws ec2 describe-instances --region {region} --filters '
                    f'Name=tag:ray-cluster-name,Values={handle.cluster_name} '
                    f'--query Reservations[].Instances[].InstanceId '
                    '--output text')
                terminate_cmd = (
                    f'aws ec2 terminate-instances --region {region} '
                    f'--instance-ids $({query_cmd})')
            elif isinstance(cloud, clouds.GCP):
                zone = config['provider']['availability_zone']
                # TODO(wei-lin): refactor by calling functions of node provider
                # that uses Python API rather than CLI
                if use_tpu_vm:
                    # check if gcloud includes TPU VM API
                    backend_utils.check_gcp_cli_include_tpu_vm()

                    query_cmd = (
                        f'gcloud compute tpus tpu-vm list --filter='
                        f'\\(labels.ray-cluster-name={cluster_name}\\) '
                        f'--zone={zone} --format=value\\(name\\)')
                    terminate_cmd = (
                        f'gcloud compute tpus tpu-vm delete --zone={zone}'
                        f' --quiet $({query_cmd})')
                else:
                    query_cmd = (
                        f'gcloud compute instances list --filter='
                        f'\\(labels.ray-cluster-name={cluster_name}\\) '
                        f'--zones={zone} --format=value\\(name\\)')
                    terminate_cmd = (
                        f'gcloud compute instances delete --zone={zone}'
                        f' --quiet $({query_cmd})')
            else:
                with ux_utils.print_exception_no_traceback():
                    raise ValueError(f'Unsupported cloud {cloud} for stopped '
                                     f'cluster {cluster_name!r}.')
            with backend_utils.safe_console_status(f'[bold cyan]Terminating '
                                                   f'[green]{cluster_name}'):
                returncode, stdout, stderr = log_lib.run_with_log(
                    terminate_cmd,
                    log_abs_path,
                    shell=True,
                    stream_logs=False,
                    require_outputs=True)
        else:
            config['provider']['cache_stopped_nodes'] = not terminate
            with tempfile.NamedTemporaryFile('w',
                                             prefix='sky_',
                                             delete=False,
                                             suffix='.yml') as f:
                common_utils.dump_yaml(f.name, config)
                f.flush()

                teardown_verb = 'Terminating' if terminate else 'Stopping'
                with backend_utils.safe_console_status(
                        f'[bold cyan]{teardown_verb} '
                        f'[green]{cluster_name}'):
                    # FIXME(zongheng): support retries. This call can fail for
                    # example due to GCP returning list requests per limit
                    # exceeded.
                    returncode, stdout, stderr = log_lib.run_with_log(
                        ['ray', 'down', '-y', f.name],
                        log_abs_path,
                        stream_logs=False,
                        require_outputs=True,
                        # Disable stdin to avoid ray outputs mess up the
                        # terminal with misaligned output when multithreading/
                        # multiprocessing are used.
                        # Refer to: https://github.com/ray-project/ray/blob/d462172be7c5779abf37609aed08af112a533e1e/python/ray/autoscaler/_private/subprocess_output_util.py#L264 # pylint: disable=line-too-long
                        stdin=subprocess.DEVNULL)
        if returncode != 0:
            if purge:
                logger.warning(
                    _TEARDOWN_PURGE_WARNING.format(
                        reason='stopping/terminating cluster nodes'))
            # This error returns when we call "gcloud delete" with an empty VM
            # list where no instance exists. Safe to ignore it and do cleanup
            # locally.
            # TODO(wei-lin): refactor error handling mechanism.
            elif 'TPU must be specified.' not in stderr:
                logger.error(
                    _TEARDOWN_FAILURE_MESSAGE.format(
                        extra_reason='',
                        cluster_name=handle.cluster_name,
                        stdout=stdout,
                        stderr=stderr))
                return False

        if post_teardown_cleanup:
            return self.post_teardown_cleanup(handle, terminate, purge)
        else:
            return True

    def post_teardown_cleanup(self,
                              handle: ResourceHandle,
                              terminate: bool,
                              purge: bool = False) -> bool:
        """Cleanup local configs/caches and delete TPUs after teardown.

        This method will handle the following cleanup steps:
        * Deleting the TPUs;
        * Removing ssh configs for the cluster;
        * Updating the local state of the cluster;
        * Removing the terminated cluster's scripts and ray yaml files.
        """
        log_path = os.path.join(os.path.expanduser(self.log_dir),
                                'teardown.log')
        log_abs_path = os.path.abspath(log_path)

        if (handle.tpu_delete_script is not None and
                os.path.exists(handle.tpu_delete_script)):
            with backend_utils.safe_console_status(
                    '[bold cyan]Terminating TPU...'):
                tpu_rc, tpu_stdout, tpu_stderr = log_lib.run_with_log(
                    ['bash', handle.tpu_delete_script],
                    log_abs_path,
                    stream_logs=False,
                    require_outputs=True)
            if tpu_rc != 0:
                if _TPU_NOT_FOUND_ERROR in tpu_stderr:
                    logger.info('TPU not found. '
                                'It should have been deleted already.')
                elif purge:
                    logger.warning(
                        _TEARDOWN_PURGE_WARNING.format(
                            reason='stopping/terminating TPU'))
                else:
                    logger.error(
                        _TEARDOWN_FAILURE_MESSAGE.format(
                            extra_reason='It is caused by TPU failure.',
                            cluster_name=handle.cluster_name,
                            stdout=tpu_stdout,
                            stderr=tpu_stderr))
                    return False

        # The cluster file must exist because the cluster_yaml will only
        # be removed after the cluster entry in the database is removed.
        config = common_utils.read_yaml(handle.cluster_yaml)
        auth_config = config['auth']
        backend_utils.SSHConfigHelper.remove_cluster(handle.cluster_name,
                                                     handle.head_ip,
                                                     auth_config)
        global_user_state.remove_cluster(handle.cluster_name,
                                         terminate=terminate)

        if terminate:
            # Clean up TPU creation/deletion scripts
            if handle.tpu_delete_script is not None:
                assert handle.tpu_create_script is not None
                os.remove(handle.tpu_create_script)
                os.remove(handle.tpu_delete_script)

            # Clean up generated config
            # No try-except is needed since Ray will fail to teardown the
            # cluster if the cluster_yaml is missing.
            os.remove(handle.cluster_yaml)
        return True

    def set_autostop(self,
                     handle: ResourceHandle,
                     idle_minutes_to_autostop: Optional[int],
                     down: bool = False,
                     stream_logs: bool = True) -> None:
        if idle_minutes_to_autostop is not None:
            code = autostop_lib.AutostopCodeGen.set_autostop(
                idle_minutes_to_autostop, self.NAME, down)
            returncode, _, stderr = self.run_on_head(handle,
                                                     code,
                                                     require_outputs=True,
                                                     stream_logs=stream_logs)
            subprocess_utils.handle_returncode(returncode,
                                               code,
                                               'Failed to set autostop',
                                               stderr=stderr,
                                               stream_logs=stream_logs)
            global_user_state.set_cluster_autostop_value(
                handle.cluster_name, idle_minutes_to_autostop, down)

    # TODO(zhwu): Refactor this to a CommandRunner class, so different backends
    # can support its own command runner.
    @timeline.event
    def run_on_head(
        self,
        handle: ResourceHandle,
        cmd: str,
        *,
        port_forward: Optional[List[str]] = None,
        log_path: str = '/dev/null',
        process_stream: bool = True,
        stream_logs: bool = False,
        use_cached_head_ip: bool = True,
        ssh_mode: command_runner.SshMode = command_runner.SshMode.
        NON_INTERACTIVE,
        under_remote_workdir: bool = False,
        require_outputs: bool = False,
        separate_stderr: bool = False,
        **kwargs,
    ) -> Union[int, Tuple[int, str, str]]:
        """Runs 'cmd' on the cluster's head node."""
        max_attempts = 1 if use_cached_head_ip else _HEAD_IP_MAX_ATTEMPTS
        head_ip = backend_utils.get_head_ip(handle, use_cached_head_ip,
                                            max_attempts)
        ssh_credentials = backend_utils.ssh_credential_from_yaml(
            handle.cluster_yaml)
        runner = command_runner.SSHCommandRunner(head_ip, **ssh_credentials)
        if under_remote_workdir:
            cmd = f'cd {SKY_REMOTE_WORKDIR} && {cmd}'

        return runner.run(
            cmd,
            port_forward=port_forward,
            log_path=log_path,
            process_stream=process_stream,
            stream_logs=stream_logs,
            ssh_mode=ssh_mode,
            require_outputs=require_outputs,
            separate_stderr=separate_stderr,
            **kwargs,
        )

    # --- Utilities ---

    @timeline.event
    def _check_existing_cluster(
            self, task: task_lib.Task, to_provision: resources_lib.Resources,
            cluster_name: str) -> RetryingVmProvisioner.ToProvisionConfig:
        handle = global_user_state.get_handle_from_cluster_name(cluster_name)
        if handle is not None:
            # Cluster already exists.
            self.check_resources_fit_cluster(handle, task)
            # Use the existing cluster.
            assert handle.launched_resources is not None, (cluster_name, handle)
            return RetryingVmProvisioner.ToProvisionConfig(
                cluster_name,
                handle.launched_resources,
                handle.launched_nodes,
                cluster_exists=True)
        usage_lib.messages.usage.set_new_cluster()
        assert len(task.resources) == 1, task.resources
        resources = list(task.resources)[0]
        task_cloud = resources.cloud
        # Use the task_cloud, because the cloud in `to_provision` can be changed
        # later during the retry.
        backend_utils.check_cluster_name_is_valid(cluster_name, task_cloud)

        cloud = to_provision.cloud
        if isinstance(cloud, clouds.Local):
            # The field ssh_user is specified in the cluster config file.
            ssh_user = onprem_utils.get_local_cluster_config_or_error(
                cluster_name)['auth']['ssh_user']
            logger.info(f'{colorama.Fore.CYAN}Connecting to local cluster: '
                        f'{cluster_name!r} [Username: {ssh_user}].'
                        f'{colorama.Style.RESET_ALL}\n'
                        'Run `sky status` to see existing clusters.')
        else:
            logger.info(
                f'{colorama.Fore.CYAN}Creating a new cluster: "{cluster_name}" '
                f'[{task.num_nodes}x {to_provision}].'
                f'{colorama.Style.RESET_ALL}\n'
                'Tip: to reuse an existing cluster, '
                'specify --cluster (-c). '
                'Run `sky status` to see existing clusters.')
        return RetryingVmProvisioner.ToProvisionConfig(cluster_name,
                                                       to_provision,
                                                       task.num_nodes)

    def _set_tpu_name(self, cluster_config_file: str, num_nodes: int,
                      tpu_name: str) -> None:
        """Sets TPU_NAME on all nodes."""
        ip_list = backend_utils.get_node_ips(cluster_config_file, num_nodes)
        ssh_credentials = backend_utils.ssh_credential_from_yaml(
            cluster_config_file)

        runners = command_runner.SSHCommandRunner.make_runner_list(
            ip_list, **ssh_credentials)

        def _setup_tpu_name_on_node(
                runner: command_runner.SSHCommandRunner) -> None:
            cmd = (f'[[ -z $TPU_NAME ]] && echo "export TPU_NAME={tpu_name}" '
                   '>> ~/.bashrc || echo "TPU_NAME already set"')
            returncode = runner.run(cmd,
                                    log_path=os.path.join(
                                        self.log_dir, 'tpu_setup.log'))
            subprocess_utils.handle_returncode(
                returncode, cmd, 'Failed to set TPU_NAME on node.')

        subprocess_utils.run_in_parallel(_setup_tpu_name_on_node, runners)

    def _execute_file_mounts(self, handle: ResourceHandle,
                             file_mounts: Dict[Path, Path]):
        """Executes file mounts - rsyncing local files and
        copying from remote stores."""
        # File mounts handling for remote paths possibly without write access:
        #  (1) in 'file_mounts' sections, add <prefix> to these target paths.
        #  (2) then, create symlinks from '/.../file' to '<prefix>/.../file'.
        if file_mounts is None or not file_mounts:
            return
        symlink_commands = []
        fore = colorama.Fore
        style = colorama.Style
        logger.info(f'{fore.CYAN}Processing file mounts.{style.RESET_ALL}')
        start = time.time()
        ip_list = backend_utils.get_node_ips(handle.cluster_yaml,
                                             handle.launched_nodes,
                                             handle=handle)
        ssh_credentials = backend_utils.ssh_credential_from_yaml(
            handle.cluster_yaml)
        runners = command_runner.SSHCommandRunner.make_runner_list(
            ip_list, **ssh_credentials)
        log_path = os.path.join(self.log_dir, 'file_mounts.log')

        # Check the files and warn
        for dst, src in file_mounts.items():
            if not data_utils.is_cloud_store_url(src):
                full_src = os.path.abspath(os.path.expanduser(src))
                # Checked during Task.set_file_mounts().
                assert os.path.exists(full_src), f'{full_src} does not exist.'
                src_size = backend_utils.path_size_megabytes(full_src)
                if src_size >= _PATH_SIZE_MEGABYTES_WARN_THRESHOLD:
                    logger.warning(
                        f'{fore.YELLOW}The size of file mount src {src!r} '
                        f'is {src_size} MB. Try to keep src small or use '
                        '.gitignore to exclude large files, as large sizes '
                        f'will slow down rsync. {style.RESET_ALL}')
                if os.path.islink(full_src):
                    logger.warning(
                        f'{fore.YELLOW}Source path {src!r} is a symlink. '
                        f'Symlink contents are not uploaded.{style.RESET_ALL}')

        tail_cmd = f'tail -n100 -f {log_path}'
        logger.info('To view detailed progress: '
                    f'{style.BRIGHT}{tail_cmd}{style.RESET_ALL}')

        for dst, src in file_mounts.items():
            # TODO: room for improvement.  Here there are many moving parts
            # (download gsutil on remote, run gsutil on remote).  Consider
            # alternatives (smart_open, each provider's own sdk), a
            # data-transfer container etc.
            if not os.path.isabs(dst) and not dst.startswith('~/'):
                dst = f'{SKY_REMOTE_WORKDIR}/{dst}'
            # Sync 'src' to 'wrapped_dst', a safe-to-write "wrapped" path.
            wrapped_dst = dst
            if not dst.startswith('~/') and not dst.startswith('/tmp/'):
                # Handles the remote paths possibly without write access.
                # (1) add <prefix> to these target paths.
                wrapped_dst = backend_utils.FileMountHelper.wrap_file_mount(dst)
                cmd = backend_utils.FileMountHelper.make_safe_symlink_command(
                    source=dst, target=wrapped_dst)
                symlink_commands.append(cmd)

            if not data_utils.is_cloud_store_url(src):
                full_src = os.path.abspath(os.path.expanduser(src))

                if os.path.isfile(full_src):
                    mkdir_for_wrapped_dst = (
                        f'mkdir -p {os.path.dirname(wrapped_dst)}')
                else:
                    mkdir_for_wrapped_dst = f'mkdir -p {wrapped_dst}'

                # TODO(mluo): Fix method so that mkdir and rsync run together
                backend_utils.parallel_data_transfer_to_nodes(
                    runners,
                    source=src,
                    target=wrapped_dst,
                    cmd=mkdir_for_wrapped_dst,
                    run_rsync=True,
                    action_message='Syncing',
                    log_path=log_path,
                    stream_logs=False,
                )
                continue

            storage = cloud_stores.get_storage_from_path(src)
            if storage.is_directory(src):
                sync = storage.make_sync_dir_command(source=src,
                                                     destination=wrapped_dst)
                # It is a directory so make sure it exists.
                mkdir_for_wrapped_dst = f'mkdir -p {wrapped_dst}'
            else:
                sync = storage.make_sync_file_command(source=src,
                                                      destination=wrapped_dst)
                # It is a file so make sure *its parent dir* exists.
                mkdir_for_wrapped_dst = \
                    f'mkdir -p {os.path.dirname(wrapped_dst)}'

            download_target_commands = [
                # Ensure sync can write to wrapped_dst (e.g., '/data/').
                mkdir_for_wrapped_dst,
                # Both the wrapped and the symlink dir exist; sync.
                sync,
            ]
            command = ' && '.join(download_target_commands)
            # dst is only used for message printing.
            backend_utils.parallel_data_transfer_to_nodes(
                runners,
                source=src,
                target=dst,
                cmd=command,
                run_rsync=False,
                action_message='Syncing',
                log_path=log_path,
                stream_logs=False,
            )
        # (2) Run the commands to create symlinks on all the nodes.
        symlink_command = ' && '.join(symlink_commands)
        if symlink_command:

            def _symlink_node(runner: command_runner.SSHCommandRunner):
                returncode = runner.run(symlink_command, log_path=log_path)
                subprocess_utils.handle_returncode(
                    returncode, symlink_command,
                    'Failed to create symlinks. The target destination '
                    f'may already exist. Log: {log_path}')

            subprocess_utils.run_in_parallel(_symlink_node, runners)
        end = time.time()
        logger.debug(f'File mount sync took {end - start} seconds.')

    def _execute_storage_mounts(self, handle: ResourceHandle,
                                storage_mounts: Dict[Path,
                                                     storage_lib.Storage]):
        """Executes storage mounts: installing mounting tools and mounting."""
        # Process only mount mode objects here. COPY mode objects have been
        # converted to regular copy file mounts and thus have been handled
        # in the '__execute_file_mounts' method.
        storage_mounts = {
            path: storage_mount
            for path, storage_mount in storage_mounts.items()
            if storage_mount.mode == storage_lib.StorageMode.MOUNT
        }

        if not storage_mounts:
            return

        cloud = handle.launched_resources.cloud
        # TODO(romil): Support Mounting for Local (remove sudo installation)
        if isinstance(cloud, clouds.Local):
            logger.warning(
                f'{colorama.Fore.YELLOW}Sky On-prem does not support '
                f'mounting. No action will be taken.{colorama.Style.RESET_ALL}')
            return

        fore = colorama.Fore
        style = colorama.Style
        plural = 's' if len(storage_mounts) > 1 else ''
        logger.info(f'{fore.CYAN}Processing {len(storage_mounts)} '
                    f'storage mount{plural}.{style.RESET_ALL}')
        start = time.time()
        ip_list = backend_utils.get_node_ips(handle.cluster_yaml,
                                             handle.launched_nodes,
                                             handle=handle)
        ssh_credentials = backend_utils.ssh_credential_from_yaml(
            handle.cluster_yaml)
        runners = command_runner.SSHCommandRunner.make_runner_list(
            ip_list, **ssh_credentials)
        log_path = os.path.join(self.log_dir, 'storage_mounts.log')

        for dst, storage_obj in storage_mounts.items():
            if not os.path.isabs(dst) and not dst.startswith('~/'):
                dst = f'{SKY_REMOTE_WORKDIR}/{dst}'
            # Get the first store and use it to mount
            store = list(storage_obj.stores.values())[0]
            mount_cmd = store.mount_command(dst)
            src_print = (storage_obj.source
                         if storage_obj.source else storage_obj.name)
            backend_utils.parallel_data_transfer_to_nodes(
                runners,
                source=src_print,
                target=dst,
                cmd=mount_cmd,
                run_rsync=False,
                action_message='Mounting',
                log_path=log_path,
            )
        end = time.time()
        logger.debug(f'Storage mount sync took {end - start} seconds.')

    def _execute_task_one_node(self, handle: ResourceHandle,
                               task: task_lib.Task, job_id: int,
                               detach_run: bool) -> None:
        # Launch the command as a Ray task.
        log_dir = os.path.join(self.log_dir, 'tasks')
        log_path = os.path.join(log_dir, 'run.log')

        accelerator_dict = backend_utils.get_task_demands_dict(task)

        codegen = RayCodeGen()
        is_local = isinstance(handle.launched_resources.cloud, clouds.Local)
        codegen.add_prologue(job_id,
                             spot_task=task.spot_task,
                             setup_cmd=self._setup_cmd,
                             envs=task.envs,
                             setup_log_path=os.path.join(log_dir, 'setup.log'),
                             is_local=is_local)
        codegen.add_gang_scheduling_placement_group(1, accelerator_dict)

        if callable(task.run):
            run_fn_code = textwrap.dedent(inspect.getsource(task.run))
            run_fn_name = task.run.__name__
            codegen.register_run_fn(run_fn_code, run_fn_name)

        # If it is a managed spot job, the JOB_ID_ENV_VAR will have been already
        # set by the controller.
        job_run_id = task.envs.get(
            constants.JOB_ID_ENV_VAR,
            common_utils.get_global_job_id(self.run_timestamp,
                                           cluster_name=handle.cluster_name,
                                           job_id=job_id))

        command_for_node = task.run if isinstance(task.run, str) else None
        use_sudo = isinstance(handle.launched_resources.cloud, clouds.Local)
        codegen.add_ray_task(
            bash_script=command_for_node,
            env_vars=task.envs,
            task_name=task.name,
            job_run_id=job_run_id,
            ray_resources_dict=backend_utils.get_task_demands_dict(task),
            log_path=log_path,
            use_sudo=use_sudo)

        codegen.add_epilogue()

        self._exec_code_on_head(handle,
                                codegen.build(),
                                job_id,
                                executable='python3',
                                detach_run=detach_run)

    def _execute_task_n_nodes(self, handle: ResourceHandle, task: task_lib.Task,
                              job_id: int, detach_run: bool) -> None:
        # Strategy:
        #   ray.init(...)
        #   for node:
        #     submit _run_cmd(cmd) with resource {node_i: 1}
        log_dir_base = self.log_dir
        log_dir = os.path.join(log_dir_base, 'tasks')
        accelerator_dict = backend_utils.get_task_demands_dict(task)

        # If TPU VM Pods is used, #num_nodes should be #num_tpu_devices
        is_tpu_vm_pod = tpu_utils.is_tpu_vm_pod(handle.launched_resources)
        if is_tpu_vm_pod:
            num_actual_nodes = tpu_utils.get_num_tpu_devices(
                handle.launched_resources)
        else:
            num_actual_nodes = task.num_nodes

        cluster_internal_ips = backend_utils.get_node_ips(handle.cluster_yaml,
                                                          handle.launched_nodes,
                                                          handle=handle,
                                                          get_internal_ips=True)
        # Ensure head node is the first element, then sort the remaining IPs
        # for stableness
        cluster_ips_sorted = [cluster_internal_ips[0]] + sorted(
            cluster_internal_ips[1:])

        codegen = RayCodeGen()
        is_local = isinstance(handle.launched_resources.cloud, clouds.Local)
        codegen.add_prologue(job_id,
                             spot_task=task.spot_task,
                             setup_cmd=self._setup_cmd,
                             envs=task.envs,
                             setup_log_path=os.path.join(log_dir, 'setup.log'),
                             is_local=is_local)
        codegen.add_gang_scheduling_placement_group(
            num_actual_nodes,
            accelerator_dict,
            cluster_ips_sorted=cluster_ips_sorted)

        if callable(task.run):
            run_fn_code = textwrap.dedent(inspect.getsource(task.run))
            run_fn_name = task.run.__name__
            codegen.register_run_fn(run_fn_code, run_fn_name)

        # If it is a managed spot job, the JOB_ID_ENV_VAR will have been already
        # set by the controller.
        job_run_id = task.envs.get(
            constants.JOB_ID_ENV_VAR,
            common_utils.get_global_job_id(self.run_timestamp,
                                           cluster_name=handle.cluster_name,
                                           job_id=job_id))

        # TODO(zhwu): The resources limitation for multi-node ray.tune and
        # horovod should be considered.
        for i in range(num_actual_nodes):
            command_for_node = task.run if isinstance(task.run, str) else None

            # Ray's per-node resources, to constrain scheduling each command to
            # the corresponding node, represented by private IPs.
            name = f'node-{i}'
            log_path = os.path.join(f'{log_dir}', f'{name}.log')
            use_sudo = isinstance(handle.launched_resources.cloud, clouds.Local)
            codegen.add_ray_task(
                bash_script=command_for_node,
                env_vars=task.envs,
                task_name=name,
                job_run_id=job_run_id,
                ray_resources_dict=accelerator_dict,
                log_path=log_path,
                gang_scheduling_id=i,
                use_sudo=use_sudo,
            )

        codegen.add_epilogue()
        # TODO(zhanghao): Add help info for downloading logs.
        self._exec_code_on_head(handle,
                                codegen.build(),
                                job_id,
                                executable='python3',
                                detach_run=detach_run)
