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

sky/provision/gcp/mig_utils.py

@@ -0,0 +1,210 @@
+"""Managed Instance Group Utils"""
+import re
+import subprocess
+from typing import Any, Dict
+
+from sky import sky_logging
+from sky.adaptors import gcp
+from sky.provision.gcp import constants
+
+logger = sky_logging.init_logger(__name__)
+
+MIG_RESOURCE_NOT_FOUND_PATTERN = re.compile(
+    r'The resource \'projects/.*/zones/.*/instanceGroupManagers/.*\' was not '
+    r'found')
+
+IT_RESOURCE_NOT_FOUND_PATTERN = re.compile(
+    r'The resource \'projects/.*/regions/.*/instanceTemplates/.*\' was not '
+    'found')
+
+
+def get_instance_template_name(cluster_name: str) -> str:
+    return f'{constants.INSTANCE_TEMPLATE_NAME_PREFIX}{cluster_name}'
+
+
+def get_managed_instance_group_name(cluster_name: str) -> str:
+    return f'{constants.MIG_NAME_PREFIX}{cluster_name}'
+
+
+def check_instance_template_exits(project_id: str, region: str,
+                                  template_name: str) -> bool:
+    compute = gcp.build('compute',
+                        'v1',
+                        credentials=None,
+                        cache_discovery=False)
+    try:
+        compute.regionInstanceTemplates().get(
+            project=project_id, region=region,
+            instanceTemplate=template_name).execute()
+    except gcp.http_error_exception() as e:
+        if IT_RESOURCE_NOT_FOUND_PATTERN.search(str(e)) is not None:
+            # Instance template does not exist.
+            return False
+        raise
+    return True
+
+
+def create_region_instance_template(cluster_name_on_cloud: str, project_id: str,
+                                    region: str, template_name: str,
+                                    node_config: Dict[str, Any]) -> dict:
+    """Create a regional instance template."""
+    logger.debug(f'Creating regional instance template {template_name!r}.')
+    compute = gcp.build('compute',
+                        'v1',
+                        credentials=None,
+                        cache_discovery=False)
+    config = node_config.copy()
+    config.pop(constants.MANAGED_INSTANCE_GROUP_CONFIG, None)
+
+    # We have to ignore user defined scheduling for DWS.
+    # TODO: Add a warning log for this behvaiour.
+    scheduling = config.get('scheduling', {})
+    assert scheduling.get('provisioningModel') != 'SPOT', (
+        'DWS does not support spot VMs.')
+
+    reservations_affinity = config.pop('reservation_affinity', None)
+    if reservations_affinity is not None:
+        logger.warning(
+            f'Ignoring reservations_affinity {reservations_affinity} '
+            'for DWS.')
+
+    # Create the regional instance template request
+    operation = compute.regionInstanceTemplates().insert(
+        project=project_id,
+        region=region,
+        body={
+            'name': template_name,
+            'properties': dict(
+                description=(
+                    'SkyPilot instance template for '
+                    f'{cluster_name_on_cloud!r} to support DWS requests.'),
+                reservationAffinity=dict(
+                    consumeReservationType='NO_RESERVATION'),
+                **config,
+            )
+        }).execute()
+    return operation
+
+
+def create_managed_instance_group(project_id: str, zone: str, group_name: str,
+                                  instance_template_url: str,
+                                  size: int) -> dict:
+    logger.debug(f'Creating managed instance group {group_name!r}.')
+    compute = gcp.build('compute',
+                        'v1',
+                        credentials=None,
+                        cache_discovery=False)
+    operation = compute.instanceGroupManagers().insert(
+        project=project_id,
+        zone=zone,
+        body={
+            'name': group_name,
+            'instanceTemplate': instance_template_url,
+            'target_size': size,
+            'instanceLifecyclePolicy': {
+                'defaultActionOnFailure': 'DO_NOTHING',
+            },
+            'updatePolicy': {
+                'type': 'OPPORTUNISTIC',
+            },
+        }).execute()
+    return operation
+
+
+def resize_managed_instance_group(project_id: str, zone: str, group_name: str,
+                                  resize_by: int, run_duration: int) -> dict:
+    logger.debug(f'Resizing managed instance group {group_name!r} by '
+                 f'{resize_by} with run duration {run_duration}.')
+    compute = gcp.build('compute',
+                        'beta',
+                        credentials=None,
+                        cache_discovery=False)
+    operation = compute.instanceGroupManagerResizeRequests().insert(
+        project=project_id,
+        zone=zone,
+        instanceGroupManager=group_name,
+        body={
+            'name': group_name,
+            'resizeBy': resize_by,
+            'requestedRunDuration': {
+                'seconds': run_duration,
+            }
+        }).execute()
+    return operation
+
+
+def cancel_all_resize_request_for_mig(project_id: str, zone: str,
+                                      group_name: str) -> None:
+    logger.debug(f'Cancelling all resize requests for MIG {group_name!r}.')
+    try:
+        compute = gcp.build('compute',
+                            'beta',
+                            credentials=None,
+                            cache_discovery=False)
+        operation = compute.instanceGroupManagerResizeRequests().list(
+            project=project_id,
+            zone=zone,
+            instanceGroupManager=group_name,
+            filter='state eq ACCEPTED').execute()
+        for request in operation.get('items', []):
+            try:
+                compute.instanceGroupManagerResizeRequests().cancel(
+                    project=project_id,
+                    zone=zone,
+                    instanceGroupManager=group_name,
+                    resizeRequest=request['name']).execute()
+            except gcp.http_error_exception() as e:
+                logger.warning('Failed to cancel resize request '
+                               f'{request["id"]!r}: {e}')
+    except gcp.http_error_exception() as e:
+        if re.search(MIG_RESOURCE_NOT_FOUND_PATTERN, str(e)) is None:
+            raise
+        logger.warning(f'MIG {group_name!r} does not exist. Skip '
+                       'resize request cancellation.')
+        logger.debug(f'Error: {e}')
+
+
+def check_managed_instance_group_exists(project_id: str, zone: str,
+                                        group_name: str) -> bool:
+    compute = gcp.build('compute',
+                        'v1',
+                        credentials=None,
+                        cache_discovery=False)
+    try:
+        compute.instanceGroupManagers().get(
+            project=project_id, zone=zone,
+            instanceGroupManager=group_name).execute()
+    except gcp.http_error_exception() as e:
+        if MIG_RESOURCE_NOT_FOUND_PATTERN.search(str(e)) is not None:
+            return False
+        raise
+    return True
+
+
+def wait_for_managed_group_to_be_stable(project_id: str, zone: str,
+                                        group_name: str, timeout: int) -> None:
+    """Wait until the managed instance group is stable."""
+    logger.debug(f'Waiting for MIG {group_name} to be stable with timeout '
+                 f'{timeout}.')
+    try:
+        cmd = ('gcloud compute instance-groups managed wait-until '
+               f'{group_name} '
+               '--stable '
+               f'--zone={zone} '
+               f'--project={project_id} '
+               f'--timeout={timeout}')
+        logger.info(
+            f'Waiting for MIG {group_name} to be stable with command:\n{cmd}')
+        proc = subprocess.run(
+            f'yes | {cmd}',
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            shell=True,
+            check=True,
+        )
+        stdout = proc.stdout.decode('ascii')
+        logger.info(stdout)
+    except subprocess.CalledProcessError as e:
+        stderr = e.stderr.decode('ascii')
+        logger.info(stderr)
+        raise

sky/provision/instance_setup.py

@@ -4,10 +4,10 @@ import functools
 import hashlib
 import json
 import os
-import resource
 import time
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
+from sky import exceptions
 from sky import provision
 from sky import sky_logging
 from sky.provision import common
@@ -15,15 +15,17 @@ from sky.provision import docker_utils
 from sky.provision import logging as provision_logging
 from sky.provision import metadata_utils
 from sky.skylet import constants
+from sky.usage import constants as usage_constants
+from sky.usage import usage_lib
 from sky.utils import accelerator_registry
 from sky.utils import command_runner
 from sky.utils import common_utils
+from sky.utils import env_options
 from sky.utils import subprocess_utils
+from sky.utils import timeline
 from sky.utils import ux_utils
 
 logger = sky_logging.init_logger(__name__)
-_START_TITLE = '\n' + '-' * 20 + 'Start: {} ' + '-' * 20
-_END_TITLE = '-' * 20 + 'End:   {} ' + '-' * 20 + '\n'
 
 _MAX_RETRY = 6
 
@@ -44,8 +46,8 @@ _RAY_PORT_COMMAND = (
     f'RAY_PORT=$({constants.SKY_PYTHON_CMD} -c '
     '"from sky.skylet import job_lib; print(job_lib.get_ray_port())" '
     '2> /dev/null || echo 6379);'
-    f'{constants.SKY_PYTHON_CMD} -c "from sky.utils import common_utils; '
-    'print(common_utils.encode_payload({\'ray_port\': $RAY_PORT}))"')
+    f'{constants.SKY_PYTHON_CMD} -c "from sky.utils import message_utils; '
+    'print(message_utils.encode_payload({\'ray_port\': $RAY_PORT}))"')
 
 # Command that calls `ray status` with SkyPilot's Ray port set.
 RAY_STATUS_WITH_SKY_RAY_PORT_COMMAND = (
@@ -68,42 +70,58 @@ MAYBE_SKYLET_RESTART_CMD = (f'{constants.ACTIVATE_SKY_REMOTE_PYTHON_ENV}; '
                             'sky.skylet.attempt_skylet;')
 
 
-def _auto_retry(func):
+def _set_usage_run_id_cmd() -> str:
+    """Gets the command to set the usage run id.
+
+    The command saves the current usage run id to the file, so that the skylet
+    can use it to report the heartbeat.
+
+    We use a function instead of a constant so that the usage run id is the
+    latest one when the function is called.
+    """
+    return (
+        f'cat {usage_constants.USAGE_RUN_ID_FILE} || '
+        # The run id is retrieved locally for the current run, so that the
+        # remote cluster will be set with the same run id as the initial
+        # launch operation.
+        f'echo "{usage_lib.messages.usage.run_id}" > '
+        f'{usage_constants.USAGE_RUN_ID_FILE}')
+
+
+def _set_skypilot_env_var_cmd() -> str:
+    """Sets the skypilot environment variables on the remote machine."""
+    env_vars = env_options.Options.all_options()
+    return '; '.join([f'export {k}={v}' for k, v in env_vars.items()])
+
+
+def _auto_retry(should_retry: Callable[[Exception], bool] = lambda _: True):
     """Decorator that retries the function if it fails.
 
     This decorator is mostly for SSH disconnection issues, which might happen
     during the setup of instances.
     """
 
-    @functools.wraps(func)
-    def retry(*args, **kwargs):
-        backoff = common_utils.Backoff(initial_backoff=1, max_backoff_factor=5)
-        for retry_cnt in range(_MAX_RETRY):
-            try:
-                return func(*args, **kwargs)
-            except Exception as e:  # pylint: disable=broad-except
-                if retry_cnt >= _MAX_RETRY - 1:
-                    raise e
-                sleep = backoff.current_backoff()
-                logger.info(
-                    f'{func.__name__}: Retrying in {sleep:.1f} seconds, '
-                    f'due to {e}')
-                time.sleep(sleep)
-
-    return retry
-
+    def decorator(func):
 
-def _log_start_end(func):
+        @functools.wraps(func)
+        def retry(*args, **kwargs):
+            backoff = common_utils.Backoff(initial_backoff=1,
+                                           max_backoff_factor=5)
+            for retry_cnt in range(_MAX_RETRY):
+                try:
+                    return func(*args, **kwargs)
+                except Exception as e:  # pylint: disable=broad-except
+                    if not should_retry(e) or retry_cnt >= _MAX_RETRY - 1:
+                        raise
+                    sleep = backoff.current_backoff()
+                    logger.info(
+                        f'{func.__name__}: Retrying in {sleep:.1f} seconds, '
+                        f'due to {e}')
+                    time.sleep(sleep)
 
-    @functools.wraps(func)
-    def wrapper(*args, **kwargs):
-        logger.info(_START_TITLE.format(func.__name__))
-        try:
-            return func(*args, **kwargs)
-        finally:
-            logger.info(_END_TITLE.format(func.__name__))
+        return retry
 
-    return wrapper
+    return decorator
 
 
 def _hint_worker_log_path(cluster_name: str, cluster_info: common.ClusterInfo,
@@ -124,7 +142,8 @@ def _parallel_ssh_with_cache(func,
     if max_workers is None:
         # Not using the default value of `max_workers` in ThreadPoolExecutor,
         # as 32 is too large for some machines.
-        max_workers = subprocess_utils.get_parallel_threads()
+        max_workers = subprocess_utils.get_parallel_threads(
+            cluster_info.provider_name)
     with futures.ThreadPoolExecutor(max_workers=max_workers) as pool:
         results = []
         runners = provision.get_command_runners(cluster_info.provider_name,
@@ -147,7 +166,7 @@ def _parallel_ssh_with_cache(func,
         return [future.result() for future in results]
 
 
-@_log_start_end
+@common.log_function_start_end
 def initialize_docker(cluster_name: str, docker_config: Dict[str, Any],
                       cluster_info: common.ClusterInfo,
                       ssh_credentials: Dict[str, Any]) -> Optional[str]:
@@ -156,7 +175,8 @@ def initialize_docker(cluster_name: str, docker_config: Dict[str, Any],
         return None
     _hint_worker_log_path(cluster_name, cluster_info, 'initialize_docker')
 
-    @_auto_retry
+    @_auto_retry(should_retry=lambda e: isinstance(e, exceptions.CommandError)
+                 and e.returncode == 255)
     def _initialize_docker(runner: command_runner.CommandRunner, log_path: str):
         docker_user = docker_utils.DockerInitializer(docker_config, runner,
                                                      log_path).initialize()
@@ -177,7 +197,8 @@ def initialize_docker(cluster_name: str, docker_config: Dict[str, Any],
     return docker_users[0]
 
 
-@_log_start_end
+@common.log_function_start_end
+@timeline.event
 def setup_runtime_on_cluster(cluster_name: str, setup_commands: List[str],
                              cluster_info: common.ClusterInfo,
                              ssh_credentials: Dict[str, Any]) -> None:
@@ -193,7 +214,7 @@ def setup_runtime_on_cluster(cluster_name: str, setup_commands: List[str],
         hasher.update(d)
     digest = hasher.hexdigest()
 
-    @_auto_retry
+    @_auto_retry()
     def _setup_node(runner: command_runner.CommandRunner, log_path: str):
         for cmd in setup_commands:
             returncode, stdout, stderr = runner.run(
@@ -253,49 +274,102 @@ def _ray_gpu_options(custom_resource: str) -> str:
     return f' --num-gpus={acc_count}'
 
 
-@_log_start_end
-@_auto_retry
-def start_ray_on_head_node(cluster_name: str, custom_resource: Optional[str],
-                           cluster_info: common.ClusterInfo,
-                           ssh_credentials: Dict[str, Any]) -> None:
-    """Start Ray on the head node."""
-    runners = provision.get_command_runners(cluster_info.provider_name,
-                                            cluster_info, **ssh_credentials)
-    head_runner = runners[0]
-    assert cluster_info.head_instance_id is not None, (cluster_name,
-                                                       cluster_info)
-
-    # Log the head node's output to the provision.log
-    log_path_abs = str(provision_logging.get_log_path())
+def ray_head_start_command(custom_resource: Optional[str],
+                           custom_ray_options: Optional[Dict[str, Any]]) -> str:
+    """Returns the command to start Ray on the head node."""
     ray_options = (
         # --disable-usage-stats in `ray start` saves 10 seconds of idle wait.
         f'--disable-usage-stats '
         f'--port={constants.SKY_REMOTE_RAY_PORT} '
         f'--dashboard-port={constants.SKY_REMOTE_RAY_DASHBOARD_PORT} '
+        f'--min-worker-port 11002 '
         f'--object-manager-port=8076 '
         f'--temp-dir={constants.SKY_REMOTE_RAY_TEMPDIR}')
     if custom_resource:
         ray_options += f' --resources=\'{custom_resource}\''
         ray_options += _ray_gpu_options(custom_resource)
+    if custom_ray_options:
+        if 'use_external_ip' in custom_ray_options:
+            custom_ray_options.pop('use_external_ip')
+        for key, value in custom_ray_options.items():
+            ray_options += f' --{key}={value}'
 
-    if cluster_info.custom_ray_options:
-        if 'use_external_ip' in cluster_info.custom_ray_options:
-            cluster_info.custom_ray_options.pop('use_external_ip')
-        for key, value in cluster_info.custom_ray_options.items():
+    cmd = (
+        f'{constants.SKY_RAY_CMD} stop; '
+        'RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 '
+        # worker_maximum_startup_concurrency controls the maximum number of
+        # workers that can be started concurrently. However, it also controls
+        # this warning message:
+        # https://github.com/ray-project/ray/blob/d5d03e6e24ae3cfafb87637ade795fb1480636e6/src/ray/raylet/worker_pool.cc#L1535-L1545
+        # maximum_startup_concurrency defaults to the number of CPUs given by
+        # multiprocessing.cpu_count() or manually specified to ray. (See
+        # https://github.com/ray-project/ray/blob/fab26e1813779eb568acba01281c6dd963c13635/python/ray/_private/services.py#L1622-L1624.)
+        # The warning will show when the number of workers is >4x the
+        # maximum_startup_concurrency, so typically 4x CPU count. However, the
+        # job controller uses 0.25cpu reservations, and each job can use two
+        # workers (one for the submitted job and one for remote actors),
+        # resulting in a worker count of 8x CPUs or more. Increase the
+        # worker_maximum_startup_concurrency to 3x CPUs so that we will only see
+        # the warning when the worker count is >12x CPUs.
+        'RAY_worker_maximum_startup_concurrency=$(( 3 * $(nproc --all) )) '
+        f'{constants.SKY_RAY_CMD} start --head {ray_options} || exit 1;' +
+        _RAY_PRLIMIT + _DUMP_RAY_PORTS + RAY_HEAD_WAIT_INITIALIZED_COMMAND)
+    return cmd
+
+
+def ray_worker_start_command(custom_resource: Optional[str],
+                             custom_ray_options: Optional[Dict[str, Any]],
+                             no_restart: bool) -> str:
+    """Returns the command to start Ray on the worker node."""
+    # We need to use the ray port in the env variable, because the head node
+    # determines the port to be used for the worker node.
+    ray_options = ('--address=${SKYPILOT_RAY_HEAD_IP}:${SKYPILOT_RAY_PORT} '
+                   '--object-manager-port=8076')
+
+    if custom_resource:
+        ray_options += f' --resources=\'{custom_resource}\''
+        ray_options += _ray_gpu_options(custom_resource)
+
+    if custom_ray_options:
+        for key, value in custom_ray_options.items():
             ray_options += f' --{key}={value}'
 
-    # Unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY to avoid using credentials
-    # from environment variables set by user. SkyPilot's ray cluster should use
-    # the `~/.aws/` credentials, as that is the one used to create the cluster,
-    # and the autoscaler module started by the `ray start` command should use
-    # the same credentials. Otherwise, `ray status` will fail to fetch the
-    # available nodes.
-    # Reference: https://github.com/skypilot-org/skypilot/issues/2441
-    cmd = (f'{constants.SKY_RAY_CMD} stop; '
-           'unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY; '
-           'RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 '
-           f'{constants.SKY_RAY_CMD} start --head {ray_options} || exit 1;' +
-           _RAY_PRLIMIT + _DUMP_RAY_PORTS + RAY_HEAD_WAIT_INITIALIZED_COMMAND)
+    cmd = (
+        'RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 '
+        f'{constants.SKY_RAY_CMD} start --disable-usage-stats {ray_options} || '
+        'exit 1;' + _RAY_PRLIMIT)
+    if no_restart:
+        # We do not use ray status to check whether ray is running, because
+        # on worker node, if the user started their own ray cluster, ray status
+        # will return 0, i.e., we don't know skypilot's ray cluster is running.
+        # Instead, we check whether the raylet process is running on gcs address
+        # that is connected to the head with the correct port.
+        cmd = (
+            f'ps aux | grep "ray/raylet/raylet" | '
+            'grep "gcs-address=${SKYPILOT_RAY_HEAD_IP}:${SKYPILOT_RAY_PORT}" '
+            f'|| {{ {cmd} }}')
+    else:
+        cmd = f'{constants.SKY_RAY_CMD} stop; ' + cmd
+    return cmd
+
+
+@common.log_function_start_end
+@_auto_retry()
+@timeline.event
+def start_ray_on_head_node(cluster_name: str, custom_resource: Optional[str],
+                           cluster_info: common.ClusterInfo,
+                           ssh_credentials: Dict[str, Any]) -> None:
+    """Start Ray on the head node."""
+    runners = provision.get_command_runners(cluster_info.provider_name,
+                                            cluster_info, **ssh_credentials)
+    head_runner = runners[0]
+    assert cluster_info.head_instance_id is not None, (cluster_name,
+                                                       cluster_info)
+
+    # Log the head node's output to the provision.log
+    log_path_abs = str(provision_logging.get_log_path())
+    cmd = ray_head_start_command(custom_resource,
+                                 cluster_info.custom_ray_options)
     logger.info(f'Running command on head node: {cmd}')
     # TODO(zhwu): add the output to log files.
     returncode, stdout, stderr = head_runner.run(
@@ -313,8 +387,9 @@ def start_ray_on_head_node(cluster_name: str, custom_resource: Optional[str],
                            f'===== stderr ====={stderr}')
 
 
-@_log_start_end
-@_auto_retry
+@common.log_function_start_end
+@_auto_retry()
+@timeline.event
 def start_ray_on_worker_nodes(cluster_name: str, no_restart: bool,
                               custom_resource: Optional[str], ray_port: int,
                               cluster_info: common.ClusterInfo,
@@ -349,43 +424,17 @@ def start_ray_on_worker_nodes(cluster_name: str, no_restart: bool,
     head_ip = (head_instance.internal_ip
                if not use_external_ip else head_instance.external_ip)
 
-    ray_options = (f'--address={head_ip}:{constants.SKY_REMOTE_RAY_PORT} '
-                   f'--object-manager-port=8076')
-
-    if custom_resource:
-        ray_options += f' --resources=\'{custom_resource}\''
-        ray_options += _ray_gpu_options(custom_resource)
-
-    if cluster_info.custom_ray_options:
-        for key, value in cluster_info.custom_ray_options.items():
-            ray_options += f' --{key}={value}'
+    ray_cmd = ray_worker_start_command(custom_resource,
+                                       cluster_info.custom_ray_options,
+                                       no_restart)
 
-    # Unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY, see the comment in
-    # `start_ray_on_head_node`.
-    cmd = (
-        f'unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY; '
-        'RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 '
-        f'{constants.SKY_RAY_CMD} start --disable-usage-stats {ray_options} || '
-        'exit 1;' + _RAY_PRLIMIT)
-    if no_restart:
-        # We do not use ray status to check whether ray is running, because
-        # on worker node, if the user started their own ray cluster, ray status
-        # will return 0, i.e., we don't know skypilot's ray cluster is running.
-        # Instead, we check whether the raylet process is running on gcs address
-        # that is connected to the head with the correct port.
-        cmd = (f'RAY_PORT={ray_port}; ps aux | grep "ray/raylet/raylet" | '
-               f'grep "gcs-address={head_ip}:${{RAY_PORT}}" || '
-               f'{{ {cmd} }}')
-    else:
-        cmd = f'{constants.SKY_RAY_CMD} stop; ' + cmd
+    cmd = (f'export SKYPILOT_RAY_HEAD_IP="{head_ip}"; '
+           f'export SKYPILOT_RAY_PORT={ray_port}; ' + ray_cmd)
 
     logger.info(f'Running command on worker nodes: {cmd}')
 
     def _setup_ray_worker(runner_and_id: Tuple[command_runner.CommandRunner,
                                                str]):
-        # for cmd in config_from_yaml['worker_start_ray_commands']:
-        #     cmd = cmd.replace('$RAY_HEAD_IP', ip_list[0][0])
-        #     runner.run(cmd)
         runner, instance_id = runner_and_id
         log_dir = metadata_utils.get_instance_log_dir(cluster_name, instance_id)
         log_path_abs = str(log_dir / ('ray_cluster' + '.log'))
@@ -398,8 +447,10 @@ def start_ray_on_worker_nodes(cluster_name: str, no_restart: bool,
             # by ray will have the correct PATH.
             source_bashrc=True)
 
+    num_threads = subprocess_utils.get_parallel_threads(
+        cluster_info.provider_name)
     results = subprocess_utils.run_in_parallel(
-        _setup_ray_worker, list(zip(worker_runners, cache_ids)))
+        _setup_ray_worker, list(zip(worker_runners, cache_ids)), num_threads)
     for returncode, stdout, stderr in results:
         if returncode:
             with ux_utils.print_exception_no_traceback():
@@ -410,8 +461,9 @@ def start_ray_on_worker_nodes(cluster_name: str, no_restart: bool,
                                    f'===== stderr ====={stderr}')
 
 
-@_log_start_end
-@_auto_retry
+@common.log_function_start_end
+@_auto_retry()
+@timeline.event
 def start_skylet_on_head_node(cluster_name: str,
                               cluster_info: common.ClusterInfo,
                               ssh_credentials: Dict[str, Any]) -> None:
@@ -425,11 +477,17 @@ def start_skylet_on_head_node(cluster_name: str,
     logger.info(f'Running command on head node: {MAYBE_SKYLET_RESTART_CMD}')
     # We need to source bashrc for skylet to make sure the autostop event can
     # access the path to the cloud CLIs.
-    returncode, stdout, stderr = head_runner.run(MAYBE_SKYLET_RESTART_CMD,
-                                                 stream_logs=False,
-                                                 require_outputs=True,
-                                                 log_path=log_path_abs,
-                                                 source_bashrc=True)
+    set_usage_run_id_cmd = _set_usage_run_id_cmd()
+    # Set the skypilot environment variables, including the usage type, debug
+    # info, and other options.
+    set_skypilot_env_var_cmd = _set_skypilot_env_var_cmd()
+    returncode, stdout, stderr = head_runner.run(
+        f'{set_usage_run_id_cmd}; {set_skypilot_env_var_cmd}; '
+        f'{MAYBE_SKYLET_RESTART_CMD}',
+        stream_logs=False,
+        require_outputs=True,
+        log_path=log_path_abs,
+        source_bashrc=True)
     if returncode:
         raise RuntimeError('Failed to start skylet on the head node '
                            f'(exit code {returncode}). Error: '
@@ -437,7 +495,7 @@ def start_skylet_on_head_node(cluster_name: str,
                            f'===== stderr ====={stderr}')
 
 
-@_auto_retry
+@_auto_retry()
 def _internal_file_mounts(file_mounts: Dict,
                           runner: command_runner.CommandRunner,
                           log_path: str) -> None:
@@ -473,28 +531,8 @@ def _internal_file_mounts(file_mounts: Dict,
         )
 
 
-def _max_workers_for_file_mounts(common_file_mounts: Dict[str, str]) -> int:
-    fd_limit, _ = resource.getrlimit(resource.RLIMIT_NOFILE)
-
-    fd_per_rsync = 5
-    for src in common_file_mounts.values():
-        if os.path.isdir(src):
-            # Assume that each file/folder under src takes 5 file descriptors
-            # on average.
-            fd_per_rsync = max(fd_per_rsync, len(os.listdir(src)) * 5)
-
-    # Reserve some file descriptors for the system and other processes
-    fd_reserve = 100
-
-    max_workers = (fd_limit - fd_reserve) // fd_per_rsync
-    # At least 1 worker, and avoid too many workers overloading the system.
-    max_workers = min(max(max_workers, 1),
-                      subprocess_utils.get_parallel_threads())
-    logger.debug(f'Using {max_workers} workers for file mounts.')
-    return max_workers
-
-
-@_log_start_end
+@common.log_function_start_end
+@timeline.event
 def internal_file_mounts(cluster_name: str, common_file_mounts: Dict[str, str],
                          cluster_info: common.ClusterInfo,
                          ssh_credentials: Dict[str, str]) -> None:
@@ -515,4 +553,5 @@ def internal_file_mounts(cluster_name: str, common_file_mounts: Dict[str, str],
         digest=None,
         cluster_info=cluster_info,
         ssh_credentials=ssh_credentials,
-        max_workers=_max_workers_for_file_mounts(common_file_mounts))
+        max_workers=subprocess_utils.get_max_workers_for_file_mounts(
+            common_file_mounts, cluster_info.provider_name))

sky/provision/kubernetes/__init__.py

@@ -2,6 +2,7 @@
 
 from sky.provision.kubernetes.config import bootstrap_instances
 from sky.provision.kubernetes.instance import get_cluster_info
+from sky.provision.kubernetes.instance import get_command_runners
 from sky.provision.kubernetes.instance import query_instances
 from sky.provision.kubernetes.instance import run_instances
 from sky.provision.kubernetes.instance import stop_instances