skypilot-nightly 1.0.0.dev20251203__py3-none-any.whl → 1.0.0.dev20260112__py3-none-any.whl

sky/server/server_utils.py

@@ -0,0 +1,30 @@
+"""Utilities for the API server."""
+
+from typing import Optional, Type, TypeVar
+
+import fastapi
+
+from sky.server.requests import payloads
+from sky.skylet import constants
+
+_BodyT = TypeVar('_BodyT', bound=payloads.RequestBody)
+
+
+# TODO(aylei): remove this and disable request body construction at server-side
+def build_body_at_server(request: Optional[fastapi.Request],
+                         body_type: Type[_BodyT], **data) -> _BodyT:
+    """Builds the request body at the server.
+
+    For historical reasons, some handlers mimic a client request body
+    at server-side in order to coordinate with the interface of executor.
+    This will cause issues where the client info like user identity is not
+    respected in these handlers. This function is a helper to build the request
+    body at server-side with the auth user overridden.
+    """
+    request_body = body_type(**data)
+    if request is not None:
+        auth_user = getattr(request.state, 'auth_user', None)
+        if auth_user:
+            request_body.env_vars[constants.USER_ID_ENV_VAR] = auth_user.id
+            request_body.env_vars[constants.USER_ENV_VAR] = auth_user.name
+    return request_body

sky/server/uvicorn.py

@@ -20,6 +20,7 @@ from uvicorn.supervisors import multiprocess
 from sky import sky_logging
 from sky.server import daemons
 from sky.server import metrics as metrics_lib
+from sky.server import plugins
 from sky.server import state
 from sky.server.requests import requests as requests_lib
 from sky.skylet import constants
@@ -237,6 +238,10 @@ def run(config: uvicorn.Config, max_db_connections: Optional[int] = None):
     server = Server(config=config, max_db_connections=max_db_connections)
     try:
         if config.workers is not None and config.workers > 1:
+            # When workers > 1, uvicorn does not run server app in the main
+            # process. In this case, plugins are not loaded at this point, so
+            # load plugins here without uvicorn app.
+            plugins.load_plugins(plugins.ExtensionContext())
             sock = config.bind_socket()
             SlowStartMultiprocess(config, target=server.run,
                                   sockets=[sock]).run()

sky/setup_files/MANIFEST.in

@@ -15,6 +15,7 @@ include sky/jobs/dashboard/templates/*
 include sky/jobs/dashboard/static/*
 include sky/templates/*
 include sky/utils/kubernetes/*
+include sky/ssh_node_pools/deploy/tunnel/*
 include sky/server/html/*
 recursive-include sky/dashboard/out *
 include sky/users/*.conf

sky/setup_files/dependencies.py

@@ -84,6 +84,7 @@ install_requires = [
     'bcrypt==4.0.1',
     'pyjwt',
     'gitpython',
+    'paramiko',
     'types-paramiko',
     'alembic',
     'aiohttp',
@@ -143,9 +144,11 @@ aws_dependencies = [
     'awscli>=1.27.10',
     'botocore>=1.29.10',
     'boto3>=1.26.1',
-    # NOTE: required by awscli. To avoid ray automatically installing
-    # the latest version.
-    'colorama < 0.4.5',
+    # NOTE: colorama is a dependency of awscli. We pin it to match the
+    # version constraint in awscli (<0.4.7) to prevent potential conflicts
+    # with other packages like ray, which might otherwise install a newer
+    # version.
+    'colorama<0.4.7',
 ]
 
 # Kubernetes 32.0.0 has an authentication bug:
@@ -203,12 +206,21 @@ cloud_dependencies: Dict[str, List[str]] = {
     'ssh': kubernetes_dependencies,
     # For the container registry auth api. Reference:
     # https://github.com/runpod/runpod-python/releases/tag/1.6.1
-    # RunPod needs a TOML parser to read ~/.runpod/config.toml. On Python 3.11+
-    # stdlib provides tomllib; on lower versions we depend on tomli explicitly.
-    # Instead of installing tomli conditionally, we install it explicitly.
-    # This is because the conditional installation of tomli does not work
-    # with controller package installation code.
-    'runpod': ['runpod>=1.6.1', 'tomli'],
+    'runpod': [
+        # For the container registry auth api. Reference:
+        # https://github.com/runpod/runpod-python/releases/tag/1.6.1
+        'runpod>=1.6.1',
+        # RunPod needs a TOML parser to read ~/.runpod/config.toml. On Python
+        # 3.11+ stdlib provides tomllib; on lower versions we depend on tomli
+        # explicitly. Instead of installing tomli conditionally, we install it
+        # explicitly. This is because the conditional installation of tomli does
+        # not work with controller package installation code.
+        'tomli',
+        # runpod installs aiodns (via aiohttp[speedups]), which is incompatible
+        # with pycares 5.0.0 due to deprecations.
+        # See https://github.com/aio-libs/aiodns/issues/214
+        'pycares<5',
+    ],
     'fluidstack': [],  # No dependencies needed for fluidstack
     'cudo': ['cudo-compute>=0.1.10'],
     'paperspace': [],  # No dependencies needed for paperspace
@@ -234,6 +246,7 @@ cloud_dependencies: Dict[str, List[str]] = {
     'hyperbolic': [],  # No dependencies needed for hyperbolic
     'seeweb': ['ecsapi==0.4.0'],
     'shadeform': [],  # No dependencies needed for shadeform
+    'slurm': ['python-hostlist'],
 }
 
 # Calculate which clouds should be included in the [all] installation.

sky/sky_logging.py

@@ -15,7 +15,8 @@ from sky.utils import env_options
 from sky.utils import rich_utils
 
 # UX: Should we show logging prefixes and some extra information in optimizer?
-_FORMAT = '%(levelname).1s %(asctime)s %(filename)s:%(lineno)d] %(message)s'
+_FORMAT = ('%(levelname).1s %(asctime)s.%(msecs)03d PID=%(process)d '
+           '%(filename)s:%(lineno)d] %(message)s')
 _DATE_FORMAT = '%m-%d %H:%M:%S'
 _SENSITIVE_LOGGER = ['sky.provisioner', 'sky.optimizer']
 

sky/skylet/attempt_skylet.py

@@ -9,6 +9,7 @@ import psutil
 
 from sky.skylet import constants
 from sky.skylet import runtime_utils
+from sky.utils import common_utils
 
 VERSION_FILE = runtime_utils.get_runtime_dir_path(constants.SKYLET_VERSION_FILE)
 SKYLET_LOG_FILE = runtime_utils.get_runtime_dir_path(constants.SKYLET_LOG_FILE)
@@ -97,8 +98,13 @@ def restart_skylet():
     for pid in _find_running_skylet_pids():
         try:
             os.kill(pid, signal.SIGKILL)
-        except (OSError, ProcessLookupError):
-            # Process died between detection and kill
+            # Wait until process fully terminates so its socket gets released.
+            # Without this, find_free_port may race with the kernel closing the
+            # socket and fail to bind to the port that's supposed to be free.
+            psutil.Process(pid).wait(timeout=5)
+        except (OSError, ProcessLookupError, psutil.NoSuchProcess,
+                psutil.TimeoutExpired):
+            # Process died between detection and kill, or timeout waiting
             pass
     # Clean up the PID file
     try:
@@ -106,7 +112,11 @@ def restart_skylet():
     except OSError:
         pass  # Best effort cleanup
 
-    port = constants.SKYLET_GRPC_PORT
+    # TODO(kevin): Handle race conditions here. Race conditions can only
+    # happen on Slurm, where there could be multiple clusters running in
+    # one network namespace. For other clouds, the behaviour will be that
+    # it always gets port 46590 (default port).
+    port = common_utils.find_free_port(constants.SKYLET_GRPC_PORT)
     subprocess.run(
         # We have made sure that `attempt_skylet.py` is executed with the
         # skypilot runtime env activated, so that skylet can access the cloud

sky/skylet/constants.py

@@ -20,11 +20,13 @@ SKY_RUNTIME_DIR = '${SKY_RUNTIME_DIR:-$HOME}'
 #    os.path.expanduser(os.environ.get(SKY_RUNTIME_DIR_ENV_VAR_KEY, '~')),
 #    '.sky/jobs.db')
 SKY_RUNTIME_DIR_ENV_VAR_KEY = 'SKY_RUNTIME_DIR'
+SKY_CLUSTER_NAME_ENV_VAR_KEY = 'SKY_CLUSTER_NAME'
 # We keep sky_logs and sky_workdir in $HOME, because
 # these are artifacts that users can access, and having
 # them be in $HOME makes it more convenient.
 SKY_LOGS_DIRECTORY = '~/sky_logs'
 SKY_REMOTE_WORKDIR = '~/sky_workdir'
+SKY_TEMPLATES_DIRECTORY = '~/sky_templates'
 SKY_IGNORE_FILE = '.skyignore'
 GIT_IGNORE_FILE = '.gitignore'
 
@@ -45,7 +47,19 @@ SKY_REMOTE_RAY_PORT_FILE = '.sky/ray_port.json'
 SKY_REMOTE_RAY_TEMPDIR = '/tmp/ray_skypilot'
 SKY_REMOTE_RAY_VERSION = '2.9.3'
 
-SKY_UNSET_PYTHONPATH = 'env -u PYTHONPATH'
+# To avoid user image causing issue with the SkyPilot runtime, we run SkyPilot
+# commands the following prefix:
+# 1. env -u PYTHONPATH: unset PYTHONPATH to avoid any package specified in
+# PYTHONPATH interfering with the SkyPilot runtime.
+# 2. env -C $HOME: set the execution directory to $HOME to avoid the case when
+# a user's WORKDIR in Dockerfile is a Python site-packages directory. Python
+# adds CWD to the beginning of sys.path, so if WORKDIR contains packages (e.g.,
+# compiled for a different Python version), imports will fail with errors like
+# "ModuleNotFoundError: No module named 'rpds.rpds'".
+#
+# TODO(zhwu): Switch -C $HOME to PYTHONSAFEPATH=1, once we moved our runtime to
+# Python 3.11 for a more robust setup.
+SKY_UNSET_PYTHONPATH_AND_SET_CWD = 'env -u PYTHONPATH -C $HOME'
 # We store the absolute path of the python executable (/opt/conda/bin/python3)
 # in this file, so that any future internal commands that need to use python
 # can use this path. This is useful for the case where the user has a custom
@@ -57,7 +71,8 @@ SKY_GET_PYTHON_PATH_CMD = (f'[ -s {SKY_PYTHON_PATH_FILE} ] && '
                            f'cat {SKY_PYTHON_PATH_FILE} 2> /dev/null || '
                            'which python3')
 # Python executable, e.g., /opt/conda/bin/python3
-SKY_PYTHON_CMD = f'{SKY_UNSET_PYTHONPATH} $({SKY_GET_PYTHON_PATH_CMD})'
+SKY_PYTHON_CMD = (f'{SKY_UNSET_PYTHONPATH_AND_SET_CWD} '
+                  f'$({SKY_GET_PYTHON_PATH_CMD})')
 # Prefer SKY_UV_PIP_CMD, which is faster.
 # TODO(cooperc): remove remaining usage (GCP TPU setup).
 SKY_PIP_CMD = f'{SKY_PYTHON_CMD} -m pip'
@@ -67,17 +82,30 @@ SKY_PIP_CMD = f'{SKY_PYTHON_CMD} -m pip'
 #   #!/opt/conda/bin/python3
 SKY_RAY_CMD = (f'{SKY_PYTHON_CMD} $([ -s {SKY_RAY_PATH_FILE} ] && '
                f'cat {SKY_RAY_PATH_FILE} 2> /dev/null || which ray)')
+
+# Use $(which env) to find env, falling back to /usr/bin/env if which is
+# unavailable. This works around a Slurm quirk where srun's execvp() doesn't
+# check execute permissions, failing when $HOME/.local/bin/env (non-executable,
+# from uv installation) shadows /usr/bin/env.
+SKY_SLURM_UNSET_PYTHONPATH = ('$(which env 2>/dev/null || echo /usr/bin/env) '
+                              '-u PYTHONPATH')
+SKY_SLURM_PYTHON_CMD = (f'{SKY_SLURM_UNSET_PYTHONPATH} '
+                        f'$({SKY_GET_PYTHON_PATH_CMD})')
+
 # Separate env for SkyPilot runtime dependencies.
 SKY_REMOTE_PYTHON_ENV_NAME = 'skypilot-runtime'
 SKY_REMOTE_PYTHON_ENV: str = f'{SKY_RUNTIME_DIR}/{SKY_REMOTE_PYTHON_ENV_NAME}'
 ACTIVATE_SKY_REMOTE_PYTHON_ENV = f'source {SKY_REMOTE_PYTHON_ENV}/bin/activate'
+# Place the conda root in the runtime directory, as installing to $HOME
+# on an NFS takes too long (1-2m slower).
+SKY_CONDA_ROOT = f'{SKY_RUNTIME_DIR}/miniconda3'
 # uv is used for venv and pip, much faster than python implementations.
 SKY_UV_INSTALL_DIR = '"$HOME/.local/bin"'
 # set UV_SYSTEM_PYTHON to false in case the
 # user provided docker image set it to true.
 # unset PYTHONPATH in case the user provided docker image set it.
 SKY_UV_CMD = ('UV_SYSTEM_PYTHON=false '
-              f'{SKY_UNSET_PYTHONPATH} {SKY_UV_INSTALL_DIR}/uv')
+              f'{SKY_UNSET_PYTHONPATH_AND_SET_CWD} {SKY_UV_INSTALL_DIR}/uv')
 # This won't reinstall uv if it's already installed, so it's safe to re-run.
 SKY_UV_INSTALL_CMD = (f'{SKY_UV_CMD} -V >/dev/null 2>&1 || '
                       'curl -LsSf https://astral.sh/uv/install.sh '
@@ -116,7 +144,7 @@ TASK_ID_LIST_ENV_VAR = f'{SKYPILOT_ENV_VAR_PREFIX}TASK_IDS'
 # cluster yaml is updated.
 #
 # TODO(zongheng,zhanghao): make the upgrading of skylet automatic?
-SKYLET_VERSION = '27'
+SKYLET_VERSION = '29'
 # The version of the lib files that skylet/jobs use. Whenever there is an API
 # change for the job_lib or log_lib, we need to bump this version, so that the
 # user can be notified to update their SkyPilot version on the remote cluster.
@@ -162,6 +190,10 @@ DISABLE_GPU_ECC_COMMAND = (
     '{ sudo reboot || echo "Failed to reboot. ECC mode may not be disabled"; } '
     '|| true; ')
 
+SETUP_SKY_DIRS_COMMANDS = (f'mkdir -p ~/sky_workdir && '
+                           f'mkdir -p ~/.sky/sky_app && '
+                           f'mkdir -p {SKY_RUNTIME_DIR}/.sky;')
+
 # Install conda on the remote cluster if it is not already installed.
 # We use conda with python 3.10 to be consistent across multiple clouds with
 # best effort.
@@ -178,8 +210,9 @@ CONDA_INSTALLATION_COMMANDS = (
     # because for some images, conda is already installed, but not initialized.
     # In this case, we need to initialize conda and set auto_activate_base to
     # true.
-    '{ bash Miniconda3-Linux.sh -b || true; '
-    'eval "$(~/miniconda3/bin/conda shell.bash hook)" && conda init && '
+    '{ '
+    f'bash Miniconda3-Linux.sh -b -p "{SKY_CONDA_ROOT}" || true; '
+    f'eval "$({SKY_CONDA_ROOT}/bin/conda shell.bash hook)" && conda init && '
     # Caller should replace {conda_auto_activate} with either true or false.
     'conda config --set auto_activate_base {conda_auto_activate} && '
     'conda activate base; }; '
@@ -222,7 +255,7 @@ _sky_version = str(version.parse(sky.__version__))
 RAY_STATUS = f'RAY_ADDRESS=127.0.0.1:{SKY_REMOTE_RAY_PORT} {SKY_RAY_CMD} status'
 RAY_INSTALLATION_COMMANDS = (
     f'{SKY_UV_INSTALL_CMD};'
-    'mkdir -p ~/sky_workdir && mkdir -p ~/.sky/sky_app;'
+    f'{SETUP_SKY_DIRS_COMMANDS}'
     # Print the PATH in provision.log to help debug PATH issues.
     'echo PATH=$PATH; '
     # Install setuptools<=69.5.1 to avoid the issue with the latest setuptools
@@ -256,7 +289,7 @@ RAY_INSTALLATION_COMMANDS = (
     #
     # Here, we add ~/.local/bin to the end of the PATH to make sure the issues
     # mentioned above are resolved.
-    'export PATH=$PATH:$HOME/.local/bin; '
+    f'export PATH=$PATH:{SKY_RUNTIME_DIR}/.local/bin; '
     # Writes ray path to file if it does not exist or the file is empty.
     f'[ -s {SKY_RAY_PATH_FILE} ] || '
     f'{{ {SKY_UV_RUN_CMD} '
@@ -264,18 +297,23 @@ RAY_INSTALLATION_COMMANDS = (
 
 # Copy SkyPilot templates from the installed wheel to ~/sky_templates.
 # This must run after the skypilot wheel is installed.
+# Note: We remove ~/sky_templates first to avoid import conflicts where Python
+# would import from ~/sky_templates instead of site-packages (because
+# sky_templates itself is a package), leading to src == dst error when
+# launching on an existing cluster.
 COPY_SKYPILOT_TEMPLATES_COMMANDS = (
+    f'rm -rf {SKY_TEMPLATES_DIRECTORY}; '
     f'{ACTIVATE_SKY_REMOTE_PYTHON_ENV}; '
     f'{SKY_PYTHON_CMD} -c \''
     'import sky_templates, shutil, os; '
     'src = os.path.dirname(sky_templates.__file__); '
-    'dst = os.path.expanduser(\"~/sky_templates\"); '
+    f'dst = os.path.expanduser(\"{SKY_TEMPLATES_DIRECTORY}\"); '
     'print(f\"Copying templates from {src} to {dst}...\"); '
-    'shutil.copytree(src, dst, dirs_exist_ok=True); '
+    'shutil.copytree(src, dst); '
     'print(f\"Templates copied successfully\")\'; '
     # Make scripts executable.
-    'find ~/sky_templates -type f ! -name "*.py" ! -name "*.md" '
-    '-exec chmod +x {} \\; ')
+    f'find {SKY_TEMPLATES_DIRECTORY} -type f ! -name "*.py" ! -name "*.md" '
+    '-exec chmod +x {} + ; ')
 
 SKYPILOT_WHEEL_INSTALLATION_COMMANDS = (
     f'{SKY_UV_INSTALL_CMD};'
@@ -438,6 +476,7 @@ OVERRIDEABLE_CONFIG_KEYS_IN_TASK: List[Tuple[str, ...]] = [
     ('gcp', 'enable_gvnic'),
     ('gcp', 'enable_gpu_direct'),
     ('gcp', 'placement_policy'),
+    ('vast', 'datacenter_only'),
     ('active_workspace',),
 ]
 # When overriding the SkyPilot configs on the API server with the client one,
@@ -498,6 +537,9 @@ SKY_USER_FILE_PATH = '~/.sky/generated'
 ENV_VAR_IS_SKYPILOT_SERVER = 'IS_SKYPILOT_SERVER'
 OVERRIDE_CONSOLIDATION_MODE = 'IS_SKYPILOT_JOB_CONTROLLER'
 IS_SKYPILOT_SERVE_CONTROLLER = 'IS_SKYPILOT_SERVE_CONTROLLER'
+# Environment variable that is set to 'true' if rolling update strategy is
+# enabled for the API server deployment.
+SKYPILOT_ROLLING_UPDATE_ENABLED = 'SKYPILOT_ROLLING_UPDATE_ENABLED'
 
 SERVE_OVERRIDE_CONCURRENT_LAUNCHES = (
     f'{SKYPILOT_ENV_VAR_PREFIX}SERVE_OVERRIDE_CONCURRENT_LAUNCHES')
@@ -532,7 +574,7 @@ CATALOG_SCHEMA_VERSION = 'v8'
 CATALOG_DIR = '~/.sky/catalogs'
 ALL_CLOUDS = ('aws', 'azure', 'gcp', 'ibm', 'lambda', 'scp', 'oci',
               'kubernetes', 'runpod', 'vast', 'vsphere', 'cudo', 'fluidstack',
-              'paperspace', 'primeintellect', 'do', 'nebius', 'ssh',
+              'paperspace', 'primeintellect', 'do', 'nebius', 'ssh', 'slurm',
               'hyperbolic', 'seeweb', 'shadeform')
 # END constants used for service catalog.
 

sky/skylet/events.py

@@ -236,7 +236,7 @@ class AutostopEvent(SkyletEvent):
                     RAY_PROVISIONER_SKYPILOT_TERMINATOR):
                 logger.info('Using new provisioner to stop the cluster.')
                 self._stop_cluster_with_new_provisioner(autostop_config, config,
-                                                        provider_name)
+                                                        provider_name, cloud)
                 return
             logger.info('Not using new provisioner to stop the cluster. '
                         f'Cloud of this cluster: {provider_name}')
@@ -314,7 +314,8 @@ class AutostopEvent(SkyletEvent):
             raise NotImplementedError
 
     def _stop_cluster_with_new_provisioner(self, autostop_config,
-                                           cluster_config, provider_name):
+                                           cluster_config, provider_name,
+                                           cloud):
         # pylint: disable=import-outside-toplevel
         from sky import provision as provision_lib
         autostop_lib.set_autostopping_started()
@@ -334,8 +335,13 @@ class AutostopEvent(SkyletEvent):
 
         # Stop the ray autoscaler to avoid scaling up, during
         # stopping/terminating of the cluster.
-        logger.info('Stopping the ray cluster.')
-        subprocess.run(f'{constants.SKY_RAY_CMD} stop', shell=True, check=True)
+        if not cloud.uses_ray():
+            logger.info('Skipping ray stop as cloud does not use Ray.')
+        else:
+            logger.info('Stopping the ray cluster.')
+            subprocess.run(f'{constants.SKY_RAY_CMD} stop',
+                           shell=True,
+                           check=True)
 
         operation_fn = provision_lib.stop_instances
         if autostop_config.down:

sky/skylet/executor/__init__.py

@@ -0,0 +1 @@
+"""Task Executors"""

sky/skylet/executor/slurm.py

@@ -0,0 +1,187 @@
+"""Slurm distributed task executor for SkyPilot.
+
+This module is invoked on each Slurm compute node via:
+    srun python -m sky.skylet.executor.slurm --script=... --log-dir=...
+"""
+import argparse
+import json
+import os
+import pathlib
+import socket
+import subprocess
+import sys
+import time
+
+import colorama
+
+from sky.skylet.log_lib import run_bash_command_with_log
+
+
+def _get_ip_address() -> str:
+    """Get the IP address of the current node."""
+    # Use socket.gethostbyname to be consistent with _get_job_node_ips(),
+    # which resolves hostnames the same way. Using `hostname -I` can return
+    # Docker bridge IPs (172.17.x.x) first, causing IP mismatch errors.
+    return socket.gethostbyname(socket.gethostname())
+
+
+def _get_job_node_ips() -> str:
+    """Get IPs of all nodes in the current Slurm job."""
+    nodelist = os.environ.get('SLURM_JOB_NODELIST', '')
+    assert nodelist, 'SLURM_JOB_NODELIST is not set'
+
+    # Expand compressed nodelist (e.g., "node[1-3,5]"
+    # -> "node1\nnode2\nnode3\nnode5")
+    result = subprocess.run(['scontrol', 'show', 'hostnames', nodelist],
+                            capture_output=True,
+                            text=True,
+                            check=False)
+    if result.returncode != 0:
+        raise RuntimeError(f'Failed to get hostnames for: {nodelist}')
+
+    hostnames = result.stdout.strip().split('\n')
+    ips = []
+    for hostname in hostnames:
+        try:
+            ip = socket.gethostbyname(hostname)
+            ips.append(ip)
+        except socket.gaierror as e:
+            raise RuntimeError('Failed to get IP for hostname: '
+                               f'{hostname}') from e
+
+    return '\n'.join(ips)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description='SkyPilot Slurm task runner for distributed execution')
+    parser.add_argument('--script', help='User script (inline, shell-quoted)')
+    parser.add_argument('--script-path',
+                        help='Path to script file (if too long for inline)')
+    parser.add_argument('--env-vars',
+                        default='{}',
+                        help='JSON-encoded environment variables')
+    parser.add_argument('--log-dir',
+                        required=True,
+                        help='Directory for log files')
+    parser.add_argument('--cluster-num-nodes',
+                        type=int,
+                        required=True,
+                        help='Total number of nodes in the cluster')
+    parser.add_argument('--cluster-ips',
+                        required=True,
+                        help='Comma-separated list of cluster node IPs')
+    parser.add_argument('--task-name',
+                        default=None,
+                        help='Task name for single-node log prefix')
+    parser.add_argument(
+        '--is-setup',
+        action='store_true',
+        help=
+        'Whether this is a setup command (affects logging prefix and filename)')
+    parser.add_argument('--alloc-signal-file',
+                        help='Path to allocation signal file')
+    parser.add_argument('--setup-done-signal-file',
+                        help='Path to setup-done signal file')
+    args = parser.parse_args()
+
+    assert args.script is not None or args.script_path is not None, (
+        'Either '
+        '--script or --script-path must be provided')
+
+    # Task rank, different from index of the node in the cluster.
+    rank = int(os.environ['SLURM_PROCID'])
+    num_nodes = int(os.environ.get('SLURM_NNODES', 1))
+    is_single_node_cluster = (args.cluster_num_nodes == 1)
+
+    # Determine node index from IP (like Ray's cluster_ips_to_node_id)
+    cluster_ips = args.cluster_ips.split(',')
+    ip_addr = _get_ip_address()
+    try:
+        node_idx = cluster_ips.index(ip_addr)
+    except ValueError as e:
+        raise RuntimeError(f'IP address {ip_addr} not found in '
+                           f'cluster IPs: {cluster_ips}') from e
+    node_name = 'head' if node_idx == 0 else f'worker{node_idx}'
+
+    # Log files are written to a shared filesystem, so each node must use a
+    # unique filename to avoid collisions.
+    if args.is_setup:
+        # TODO(kevin): This is inconsistent with other clouds, where it is
+        # simply called 'setup.log'. On Slurm that is obviously not possible,
+        # since the ~/sky_logs directory is shared by all nodes, so
+        # 'setup.log' will be overwritten by other nodes.
+        # Perhaps we should apply this naming convention to other clouds.
+        log_filename = f'setup-{node_name}.log'
+    elif is_single_node_cluster:
+        log_filename = 'run.log'
+    else:
+        log_filename = f'{rank}-{node_name}.log'
+    log_path = os.path.join(args.log_dir, log_filename)
+
+    if args.script_path:
+        with open(args.script_path, 'r', encoding='utf-8') as f:
+            script = f.read()
+    else:
+        script = args.script
+
+    # Parse env vars and add SKYPILOT environment variables
+    env_vars = json.loads(args.env_vars)
+    if not args.is_setup:
+        # For setup, env vars are set in CloudVmRayBackend._setup.
+        env_vars['SKYPILOT_NODE_RANK'] = str(rank)
+        env_vars['SKYPILOT_NUM_NODES'] = str(num_nodes)
+        env_vars['SKYPILOT_NODE_IPS'] = _get_job_node_ips()
+
+    # Signal file coordination for setup/run synchronization
+    # Rank 0 touches the allocation signal to indicate resources acquired
+    if args.alloc_signal_file is not None and rank == 0:
+        pathlib.Path(args.alloc_signal_file).touch()
+
+    # Wait for setup to complete.
+    while args.setup_done_signal_file is not None and not os.path.exists(
+            args.setup_done_signal_file):
+        time.sleep(0.1)
+
+    # Build log prefix
+    # For setup on head: (setup pid={pid})
+    # For setup on workers: (setup pid={pid}, ip=1.2.3.4)
+    # For single-node cluster: (task_name, pid={pid})
+    # For multi-node on head: (head, rank=0, pid={pid})
+    # For multi-node on workers: (worker1, rank=1, pid={pid}, ip=1.2.3.4)
+    # The {pid} placeholder will be replaced by run_with_log
+    if args.is_setup:
+        # Setup prefix: head (node_idx=0) shows no IP, workers show IP
+        if node_idx == 0:
+            prefix = (f'{colorama.Fore.CYAN}(setup pid={{pid}})'
+                      f'{colorama.Style.RESET_ALL} ')
+        else:
+            prefix = (f'{colorama.Fore.CYAN}(setup pid={{pid}}, ip={ip_addr})'
+                      f'{colorama.Style.RESET_ALL} ')
+    elif is_single_node_cluster:
+        # Single-node cluster: use task name
+        name_str = args.task_name if args.task_name else 'task'
+        prefix = (f'{colorama.Fore.CYAN}({name_str}, pid={{pid}})'
+                  f'{colorama.Style.RESET_ALL} ')
+    else:
+        # Multi-node cluster: head (node_idx=0) shows no IP, workers show IP
+        if node_idx == 0:
+            prefix = (
+                f'{colorama.Fore.CYAN}({node_name}, rank={rank}, pid={{pid}})'
+                f'{colorama.Style.RESET_ALL} ')
+        else:
+            prefix = (f'{colorama.Fore.CYAN}'
+                      f'({node_name}, rank={rank}, pid={{pid}}, ip={ip_addr})'
+                      f'{colorama.Style.RESET_ALL} ')
+
+    returncode = run_bash_command_with_log(script,
+                                           log_path,
+                                           env_vars=env_vars,
+                                           stream_logs=True,
+                                           streaming_prefix=prefix)
+
+    sys.exit(returncode)
+
+
+if __name__ == '__main__':
+    main()