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

sky/utils/command_runner.py

@@ -1,14 +1,23 @@
 """Runner for commands to be executed on the cluster."""
 import enum
+import fcntl
 import hashlib
 import os
 import pathlib
+import pty
 import re
 import shlex
+import signal
+import socket
 import sys
+import termios
+import threading
 import time
 from typing import (Any, Callable, Dict, Iterable, List, Optional, Tuple, Type,
                     Union)
+import uuid
+
+import colorama
 
 from sky import exceptions
 from sky import sky_logging
@@ -19,6 +28,7 @@ from sky.utils import common_utils
 from sky.utils import context_utils
 from sky.utils import control_master_utils
 from sky.utils import git as git_utils
+from sky.utils import interactive_utils
 from sky.utils import subprocess_utils
 from sky.utils import timeline
 
@@ -63,6 +73,22 @@ def _ssh_control_path(ssh_control_filename: Optional[str]) -> Optional[str]:
     return path
 
 
+def _is_skypilot_managed_key(key_path: str) -> bool:
+    """Check if SSH key follows SkyPilot's managed key format.
+
+    SkyPilot-managed keys follow the pattern: ~/.sky/clients/<hash>/ssh/sky-key
+    External keys (like ~/.ssh/id_rsa) do not follow this pattern.
+
+    Args:
+        key_path: Path to the SSH private key.
+
+    Returns:
+        True if the key follows SkyPilot's managed format, False otherwise.
+    """
+    parts = os.path.normpath(key_path).split(os.path.sep)
+    return len(parts) >= 2 and parts[-1] == 'sky-key' and parts[-2] == 'ssh'
+
+
 # Disable sudo for root user. This is useful when the command is running in a
 # docker container, i.e. image_id is a docker image.
 ALIAS_SUDO_TO_EMPTY_FOR_ROOT_CMD = (
@@ -74,10 +100,12 @@ def ssh_options_list(
     ssh_control_name: Optional[str],
     *,
     ssh_proxy_command: Optional[str] = None,
+    ssh_proxy_jump: Optional[str] = None,
     docker_ssh_proxy_command: Optional[str] = None,
     connect_timeout: Optional[int] = None,
     port: int = 22,
     disable_control_master: Optional[bool] = False,
+    escape_percent_expand: bool = False,
 ) -> List[str]:
     """Returns a list of sane options for 'ssh'."""
     if connect_timeout is None:
@@ -117,11 +145,11 @@ def ssh_options_list(
     # SSH Control will have a severe delay when using docker_ssh_proxy_command.
     # TODO(tian): Investigate why.
     #
-    # We disable ControlMaster when ssh_proxy_command is used, because the
-    # master connection will be idle although the connection might be shared
-    # by other ssh commands that is not idle. In that case, user's custom proxy
-    # command may drop the connection due to idle timeout, since it will only
-    # see the idle master connection. It is an issue even with the
+    # We disable ControlMaster when ssh_proxy_command is used,
+    # because the master connection will be idle although the connection might
+    # be shared by other ssh commands that is not idle. In that case, user's
+    # custom proxy command may drop the connection due to idle timeout, since it
+    # will only see the idle master connection. It is an issue even with the
     # ServerAliveInterval set, since the keepalive message may not be recognized
     # by the custom proxy command, such as AWS SSM Session Manager.
     #
@@ -132,11 +160,14 @@ def ssh_options_list(
     # 'ControlPersist' number of seconds delay per ssh commands ran.
     if (ssh_control_name is not None and docker_ssh_proxy_command is None and
             ssh_proxy_command is None and not disable_control_master):
+        control_path = f'{_ssh_control_path(ssh_control_name)}/%C'
+        if escape_percent_expand:
+            control_path = control_path.replace('%', '%%')
         arg_dict.update({
             # Control path: important optimization as we do multiple ssh in one
             # sky.launch().
             'ControlMaster': 'auto',
-            'ControlPath': f'{_ssh_control_path(ssh_control_name)}/%C',
+            'ControlPath': control_path,
             'ControlPersist': '300s',
         })
     ssh_key_option = [
@@ -158,6 +189,15 @@ def ssh_options_list(
             'ProxyCommand': shlex.quote(ssh_proxy_command),
         })
 
+    if ssh_proxy_jump is not None:
+        logger.debug(f'--- ProxyJump: {ssh_proxy_jump} ---')
+        if ssh_proxy_command is not None:
+            logger.warning('Both ProxyCommand and ProxyJump are specified. '
+                           'ProxyCommand will take precedence.')
+        arg_dict.update({
+            'ProxyJump': shlex.quote(ssh_proxy_jump),
+        })
+
     return ssh_key_option + [
         x for y in (['-o', f'{k}={v}']
                     for k, v in arg_dict.items()
@@ -217,6 +257,7 @@ class CommandRunner:
         skip_num_lines: int,
         source_bashrc: bool = False,
         use_login: bool = True,
+        run_in_background: bool = False,
     ) -> str:
         """Returns the command to run."""
         if isinstance(cmd, list):
@@ -247,7 +288,11 @@ class CommandRunner:
             ]
         if not separate_stderr:
             command.append('2>&1')
+        if run_in_background:
+            command = ['nohup'] + command + ['&']
         if not process_stream and skip_num_lines:
+            assert not run_in_background, (
+                'run_in_background and skip_num_lines cannot be used together')
             command += [
                 # A hack to remove the following bash warnings (twice):
                 #  bash: cannot set terminal process group
@@ -408,6 +453,7 @@ class CommandRunner:
             connect_timeout: Optional[int] = None,
             source_bashrc: bool = False,
             skip_num_lines: int = 0,
+            run_in_background: bool = False,
             **kwargs) -> Union[int, Tuple[int, str, str]]:
         """Runs the command on the cluster.
 
@@ -426,6 +472,7 @@ class CommandRunner:
                 output. This is used when the output is not processed by
                 SkyPilot but we still want to get rid of some warning messages,
                 such as SSH warnings.
+            run_in_background: Whether to run the command in the background.
 
         Returns:
             returncode
@@ -603,17 +650,19 @@ class SSHCommandRunner(CommandRunner):
         self,
         node: Tuple[str, int],
         ssh_user: str,
-        ssh_private_key: str,
+        ssh_private_key: Optional[str],
         ssh_control_name: Optional[str] = '__default__',
         ssh_proxy_command: Optional[str] = None,
+        ssh_proxy_jump: Optional[str] = None,
         docker_user: Optional[str] = None,
         disable_control_master: Optional[bool] = False,
         port_forward_execute_remote_command: Optional[bool] = False,
+        enable_interactive_auth: bool = False,
     ):
         """Initialize SSHCommandRunner.
 
         Example Usage:
-            runner = SSHCommandRunner(ip, ssh_user, ssh_private_key)
+            runner = SSHCommandRunner((ip, port), ssh_user, ssh_private_key)
             runner.run('ls -l', mode=SshMode.NON_INTERACTIVE)
             runner.rsync(source, target, up=True)
 
@@ -628,6 +677,8 @@ class SSHCommandRunner(CommandRunner):
             ssh_proxy_command: Optional, the value to pass to '-o
                 ProxyCommand'. Useful for communicating with clusters without
                 public IPs using a "jump server".
+            ssh_proxy_jump: Optional, the value to pass to '-o ProxyJump' flag.
+                Similar to ssh_proxy_command, but more modern.
             port: The port to use for ssh.
             docker_user: The docker user to use for ssh. If specified, the
                 command will be run inside a docker container which have a ssh
@@ -647,11 +698,21 @@ class SSHCommandRunner(CommandRunner):
             None if ssh_control_name is None else hashlib.md5(
                 ssh_control_name.encode()).hexdigest()[:_HASH_MAX_LENGTH])
         self._ssh_proxy_command = ssh_proxy_command
+        self._ssh_proxy_jump = ssh_proxy_jump
         self.disable_control_master = (
             disable_control_master or
             control_master_utils.should_disable_control_master())
-        # ensure the ssh key files are created from the database
-        auth_utils.create_ssh_key_files_from_db(ssh_private_key)
+        # Ensure SSH key is available. For SkyPilot-managed keys, create from
+        # database. For external keys (e.g., Slurm clusters), verify existence.
+        if ssh_private_key is not None and _is_skypilot_managed_key(
+                ssh_private_key):
+            auth_utils.create_ssh_key_files_from_db(ssh_private_key)
+        elif ssh_private_key is not None:
+            # Externally managed key - just verify it exists
+            expanded_key_path = os.path.expanduser(ssh_private_key)
+            if not os.path.exists(expanded_key_path):
+                raise FileNotFoundError(
+                    f'SSH private key not found: {expanded_key_path}')
         if docker_user is not None:
             assert port is None or port == 22, (
                 f'port must be None or 22 for docker_user, got {port}.')
@@ -687,6 +748,7 @@ class SSHCommandRunner(CommandRunner):
             self._docker_ssh_proxy_command = None
         self.port_forward_execute_remote_command = (
             port_forward_execute_remote_command)
+        self.enable_interactive_auth = enable_interactive_auth
 
     def port_forward_command(
             self,
@@ -738,6 +800,7 @@ class SSHCommandRunner(CommandRunner):
             self.ssh_private_key,
             self.ssh_control_name,
             ssh_proxy_command=self._ssh_proxy_command,
+            ssh_proxy_jump=self._ssh_proxy_jump,
             docker_ssh_proxy_command=docker_ssh_proxy_command,
             port=self.port,
             connect_timeout=connect_timeout,
@@ -745,6 +808,127 @@ class SSHCommandRunner(CommandRunner):
                 f'{self.ssh_user}@{self.ip}'
             ]
 
+    def _retry_with_interactive_auth(
+            self, session_id: str, command: List[str], log_path: str,
+            require_outputs: bool, process_stream: bool, stream_logs: bool,
+            executable: str,
+            **kwargs) -> Union[int, Tuple[int, str, str], Tuple[int, int]]:
+        """Retries command with interactive auth.
+
+        This handles SSH connections requiring keyboard-interactive
+        authentication (e.g., 2FA) by using a PTY for auth prompts and
+        establishing a persistent ControlMaster socket (if enabled) that
+        other SSH sessions can reuse without re-authenticating.
+
+        The PTY is bridged to a websocket connection that allows the client
+        to handle interactive authentication. Command output flows through
+        normal stdout/stderr pipes, which gets printed to log_path.
+
+        See ssh_options_list for when ControlMaster is not enabled.
+        """
+        extra_options = [
+            # Override ControlPersist to reduce frequency of manual user
+            # intervention. The default from ssh_options_list is only 5m.
+            #
+            # NOTE: When used with ProxyJump, the connection can die
+            # earlier than expected, so it is recommended to also enable
+            # ControlMaster on the jump host's SSH config. It is hard to
+            # tell why exactly, because enabling -v makes this problem
+            # disappear for some reasons.
+            '-o',
+            'ControlPersist=1d',
+        ]
+        if self._ssh_proxy_jump is not None:
+            logger.warning(f'{colorama.Fore.YELLOW}When using ProxyJump, it is '
+                           'recommended to also enable ControlMaster on the '
+                           'jump host\'s SSH config to keep the authenticated '
+                           f'connection alive for longer.{colorama.Fore.RESET}')
+        command = command[:1] + extra_options + command[1:]
+
+        # Create PTY for SSH. PTY slave for stdin from user, PTY master
+        # for password/auth prompts from SSH.
+        pty_m_fd, pty_s_fd = pty.openpty()
+
+        # Create Unix socket to pass PTY master fd to websocket handler
+        fd_socket_path = interactive_utils.get_pty_socket_path(session_id)
+        if os.path.exists(fd_socket_path):
+            os.unlink(fd_socket_path)
+        fd_server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+        fd_server.bind(fd_socket_path)
+        fd_server.listen(1)
+        fd_server.settimeout(60)
+
+        # Signal client to initiate websocket for interactive auth
+        interactive_signal = f'<sky-interactive session="{session_id}"/>'
+        print(interactive_signal, flush=True)
+
+        def handle_unix_socket_connection():
+            """Background thread to handle Unix socket connection."""
+            conn = None
+            try:
+                # Wait for websocket handler to connect.
+                conn, _ = fd_server.accept()
+                # Send PTY master fd through Unix socket.
+                interactive_utils.send_fd(conn, pty_m_fd)
+                # We don't need to block here to wait for the websocket
+                # handler, as SSH will continue by itself once auth
+                # is complete.
+            except socket.timeout:
+                logger.debug('Timeout waiting for interactive auth connection')
+            except Exception as e:  # pylint: disable=broad-except
+                logger.error(f'Error in Unix socket connection: '
+                             f'{common_utils.format_exception(e)}')
+            finally:
+                if conn is not None:
+                    try:
+                        conn.close()
+                    except Exception:  # pylint: disable=broad-except
+                        pass
+                try:
+                    os.close(pty_m_fd)
+                except Exception:  # pylint: disable=broad-except
+                    pass
+
+        unix_sock_thread = threading.Thread(
+            target=handle_unix_socket_connection, daemon=True)
+        unix_sock_thread.start()
+
+        try:
+
+            def setup_pty_session():
+                # Set PTY as controlling terminal so SSH can access /dev/tty
+                # for keyboard-interactive auth. Without this:
+                # "can't open /dev/tty: Device not configured"
+                fcntl.ioctl(pty_s_fd, termios.TIOCSCTTY, 0)
+                # Ignore SIGHUP so ControlMaster survives when PTY closes.
+                signal.signal(signal.SIGHUP, signal.SIG_IGN)
+                # Ignore SIGTERM so ControlMaster survives subprocess_daemon
+                # killing the process group.
+                if self._ssh_proxy_jump is not None:
+                    signal.signal(signal.SIGTERM, signal.SIG_IGN)
+
+            return log_lib.run_with_log(' '.join(command),
+                                        log_path,
+                                        require_outputs=require_outputs,
+                                        stream_logs=stream_logs,
+                                        process_stream=process_stream,
+                                        shell=True,
+                                        executable=executable,
+                                        preexec_fn=setup_pty_session,
+                                        **kwargs)
+        except Exception as e:
+            raise RuntimeError(f'Exception in setup: {e}') from e
+        finally:
+            # Clean up PTY fds and sockets.
+            fd_server.close()
+            if os.path.exists(fd_socket_path):
+                os.unlink(fd_socket_path)
+            try:
+                os.close(pty_m_fd)
+            except OSError:
+                pass  # Already closed by background thread
+            os.close(pty_s_fd)
+
     def close_cached_connection(self) -> None:
         """Close the cached connection to the remote machine.
 
@@ -785,6 +969,7 @@ class SSHCommandRunner(CommandRunner):
             connect_timeout: Optional[int] = None,
             source_bashrc: bool = False,
             skip_num_lines: int = 0,
+            run_in_background: bool = False,
             **kwargs) -> Union[int, Tuple[int, str, str]]:
         """Uses 'ssh' to run 'cmd' on a node with ip.
 
@@ -809,27 +994,32 @@ class SSHCommandRunner(CommandRunner):
                 output. This is used when the output is not processed by
                 SkyPilot but we still want to get rid of some warning messages,
                 such as SSH warnings.
+            run_in_background: Whether to run the command in the background.
 
         Returns:
             returncode
             or
             A tuple of (returncode, stdout, stderr).
         """
+
         base_ssh_command = self.ssh_base_command(
             ssh_mode=ssh_mode,
             port_forward=port_forward,
             connect_timeout=connect_timeout)
+
         if ssh_mode == SshMode.LOGIN:
             assert isinstance(cmd, list), 'cmd must be a list for login mode.'
             command = base_ssh_command + cmd
             proc = subprocess_utils.run(command, shell=False, check=False)
             return proc.returncode, '', ''
 
-        command_str = self._get_command_to_run(cmd,
-                                               process_stream,
-                                               separate_stderr,
-                                               skip_num_lines=skip_num_lines,
-                                               source_bashrc=source_bashrc)
+        command_str = self._get_command_to_run(
+            cmd,
+            process_stream,
+            separate_stderr,
+            skip_num_lines=skip_num_lines,
+            source_bashrc=source_bashrc,
+            run_in_background=run_in_background)
         command = base_ssh_command + [shlex.quote(command_str)]
 
         log_dir = os.path.expanduser(os.path.dirname(log_path))
@@ -847,14 +1037,35 @@ class SSHCommandRunner(CommandRunner):
             else:
                 command += [f'> {log_path}']
             executable = '/bin/bash'
-        return log_lib.run_with_log(' '.join(command),
-                                    log_path,
-                                    require_outputs=require_outputs,
-                                    stream_logs=stream_logs,
-                                    process_stream=process_stream,
-                                    shell=True,
-                                    executable=executable,
-                                    **kwargs)
+
+        result = log_lib.run_with_log(' '.join(command),
+                                      log_path,
+                                      require_outputs=require_outputs,
+                                      stream_logs=stream_logs,
+                                      process_stream=process_stream,
+                                      shell=True,
+                                      executable=executable,
+                                      **kwargs)
+        if not self.enable_interactive_auth:
+            return result
+
+        if require_outputs:
+            returncode, _, _ = result
+        else:
+            returncode = result
+
+        if returncode != 255:
+            return result
+        # Exit code 255 indicates an SSH connection error. It does not
+        # necessarily mean an auth failure, but when ControlMaster is used,
+        # the stdout/stderr does not contain the auth failure message,
+        # which is why we don't check the output here, and just attempt
+        # the interactive auth flow.
+        session_id = str(uuid.uuid4())
+        return self._retry_with_interactive_auth(session_id, command, log_path,
+                                                 require_outputs,
+                                                 process_stream, stream_logs,
+                                                 executable, **kwargs)
 
     @timeline.event
     def rsync(
@@ -867,6 +1078,7 @@ class SSHCommandRunner(CommandRunner):
         log_path: str = os.devnull,
         stream_logs: bool = True,
         max_retry: int = 1,
+        get_remote_home_dir: Callable[[], str] = lambda: '~',
     ) -> None:
         """Uses 'rsync' to sync 'source' to 'target'.
 
@@ -879,6 +1091,8 @@ class SSHCommandRunner(CommandRunner):
             stream_logs: Stream logs to the stdout/stderr.
             max_retry: The maximum number of retries for the rsync command.
               This value should be non-negative.
+            get_remote_home_dir: A callable that returns the remote home
+              directory. Defaults to '~'.
 
         Raises:
             exceptions.CommandError: rsync command failed.
@@ -892,6 +1106,7 @@ class SSHCommandRunner(CommandRunner):
                 self.ssh_private_key,
                 self.ssh_control_name,
                 ssh_proxy_command=self._ssh_proxy_command,
+                ssh_proxy_jump=self._ssh_proxy_jump,
                 docker_ssh_proxy_command=docker_ssh_proxy_command,
                 port=self.port,
                 disable_control_master=self.disable_control_master))
@@ -903,7 +1118,8 @@ class SSHCommandRunner(CommandRunner):
                     rsh_option=rsh_option,
                     log_path=log_path,
                     stream_logs=stream_logs,
-                    max_retry=max_retry)
+                    max_retry=max_retry,
+                    get_remote_home_dir=get_remote_home_dir)
 
 
 class KubernetesCommandRunner(CommandRunner):
@@ -1004,6 +1220,7 @@ class KubernetesCommandRunner(CommandRunner):
             connect_timeout: Optional[int] = None,
             source_bashrc: bool = False,
             skip_num_lines: int = 0,
+            run_in_background: bool = False,
             **kwargs) -> Union[int, Tuple[int, str, str]]:
         """Uses 'kubectl exec' to run 'cmd' on a pod or deployment by its
         name and namespace.
@@ -1028,6 +1245,7 @@ class KubernetesCommandRunner(CommandRunner):
                 output. This is used when the output is not processed by
                 SkyPilot but we still want to get rid of some warning messages,
                 such as SSH warnings.
+            run_in_background: Whether to run the command in the background.
 
         Returns:
             returncode
@@ -1064,11 +1282,13 @@ class KubernetesCommandRunner(CommandRunner):
             kubectl_base_command.append('-i')
         kubectl_base_command += [*kubectl_args, '--']
 
-        command_str = self._get_command_to_run(cmd,
-                                               process_stream,
-                                               separate_stderr,
-                                               skip_num_lines=skip_num_lines,
-                                               source_bashrc=source_bashrc)
+        command_str = self._get_command_to_run(
+            cmd,
+            process_stream,
+            separate_stderr,
+            skip_num_lines=skip_num_lines,
+            source_bashrc=source_bashrc,
+            run_in_background=run_in_background)
         command = kubectl_base_command + [
             # It is important to use /bin/bash -c here to make sure we quote the
             # command to be run properly. Otherwise, directly appending commands
@@ -1182,16 +1402,19 @@ class LocalProcessCommandRunner(CommandRunner):
             connect_timeout: Optional[int] = None,
             source_bashrc: bool = False,
             skip_num_lines: int = 0,
+            run_in_background: bool = False,
             **kwargs) -> Union[int, Tuple[int, str, str]]:
         """Use subprocess to run the command."""
         del port_forward, ssh_mode, connect_timeout  # Unused.
 
-        command_str = self._get_command_to_run(cmd,
-                                               process_stream,
-                                               separate_stderr,
-                                               skip_num_lines=skip_num_lines,
-                                               source_bashrc=source_bashrc,
-                                               use_login=False)
+        command_str = self._get_command_to_run(
+            cmd,
+            process_stream,
+            separate_stderr,
+            skip_num_lines=skip_num_lines,
+            source_bashrc=source_bashrc,
+            use_login=False,
+            run_in_background=run_in_background)
 
         log_dir = os.path.expanduser(os.path.dirname(log_path))
         os.makedirs(log_dir, exist_ok=True)
@@ -1247,3 +1470,134 @@ class LocalProcessCommandRunner(CommandRunner):
                     log_path=log_path,
                     stream_logs=stream_logs,
                     max_retry=max_retry)
+
+
+class SlurmCommandRunner(SSHCommandRunner):
+    """Runner for Slurm commands.
+
+    SlurmCommandRunner sends commands over an SSH connection through the Slurm
+    controller, to the virtual instances.
+    """
+
+    def __init__(
+        self,
+        node: Tuple[str, int],
+        ssh_user: str,
+        ssh_private_key: Optional[str],
+        *,
+        sky_dir: str,
+        skypilot_runtime_dir: str,
+        job_id: str,
+        slurm_node: str,
+        **kwargs,
+    ):
+        """Initialize SlurmCommandRunner.
+
+        Example Usage:
+            runner = SlurmCommandRunner(
+                (ip, port),
+                ssh_user,
+                ssh_private_key,
+                sky_dir=sky_dir,
+                skypilot_runtime_dir=skypilot_runtime_dir,
+                job_id=job_id,
+                slurm_node=slurm_node)
+            runner.run('ls -l', mode=SshMode.NON_INTERACTIVE)
+            runner.rsync(source, target, up=True)
+
+        Args:
+            node: (ip, port) The IP address and port of the remote machine
+              (login node).
+            ssh_user: SSH username.
+            ssh_private_key: Path to SSH private key.
+            sky_dir: The private directory for the SkyPilot cluster on the
+              Slurm cluster.
+            skypilot_runtime_dir: The directory for the SkyPilot runtime
+              on the Slurm cluster.
+            job_id: The Slurm job ID for this instance.
+            slurm_node: The Slurm node hostname for this instance
+              (compute node).
+            **kwargs: Additional arguments forwarded to SSHCommandRunner
+              (e.g., ssh_proxy_command).
+        """
+        super().__init__(node, ssh_user, ssh_private_key, **kwargs)
+        self.sky_dir = sky_dir
+        self.skypilot_runtime_dir = skypilot_runtime_dir
+        self.job_id = job_id
+        self.slurm_node = slurm_node
+
+    def rsync(
+        self,
+        source: str,
+        target: str,
+        *,
+        up: bool,
+        log_path: str = os.devnull,
+        stream_logs: bool = True,
+        max_retry: int = 1,
+    ) -> None:
+        """Rsyncs files to/from the Slurm compute node using srun as transport.
+        """
+        ssh_command = ' '.join(
+            self.ssh_base_command(ssh_mode=SshMode.NON_INTERACTIVE,
+                                  port_forward=None,
+                                  connect_timeout=None))
+
+        # rsh command: parse job_id+node_list from $1, ssh to login node,
+        # run srun with rsync command.
+        rsh_option = (
+            f'bash --norc --noprofile -c \''
+            f'job_id=$(echo "$1" | cut -d+ -f1); '
+            f'node_list=$(echo "$1" | cut -d+ -f2); '
+            f'shift; '  # Shift past the encoded job_id+node_list
+            f'exec {ssh_command} '  # SSH to login node to run srun
+            f'srun --unbuffered --quiet --overlap '
+            f'--jobid="$job_id" --nodelist="$node_list" --nodes=1 --ntasks=1 '
+            f'"$@"'
+            f'\' --')
+        encoded_info = f'{self.job_id}+{self.slurm_node}'
+        self._rsync(source,
+                    target,
+                    node_destination=encoded_info,
+                    up=up,
+                    rsh_option=rsh_option,
+                    log_path=log_path,
+                    stream_logs=stream_logs,
+                    max_retry=max_retry,
+                    get_remote_home_dir=lambda: self.sky_dir)
+
+    @timeline.event
+    @context_utils.cancellation_guard
+    def run(self, cmd: Union[str, List[str]],
+            **kwargs) -> Union[int, Tuple[int, str, str]]:
+        """Run Slurm-supported user commands over an SSH connection.
+
+        Args:
+            cmd: The Slurm-supported user command to run.
+
+        Returns:
+            returncode
+            or
+            A tuple of (returncode, stdout, stderr).
+        """
+        # Override $HOME so that each SkyPilot cluster's state is isolated
+        # from one another. We rely on the assumption that ~ is exclusively
+        # used by a cluster, and in Slurm that is not the case, as $HOME
+        # could be part of a shared filesystem.
+        # And similarly for SKY_RUNTIME_DIR. See constants.\
+        # SKY_RUNTIME_DIR_ENV_VAR_KEY for more details.
+        cmd = (
+            f'export {constants.SKY_RUNTIME_DIR_ENV_VAR_KEY}='
+            f'"{self.skypilot_runtime_dir}" && '
+            # Set the uv cache directory to /tmp/uv_cache_$(id -u) to speed up
+            # package installation while avoiding permission conflicts when
+            # multiple users share the same host. Otherwise it defaults to
+            # ~/.cache/uv.
+            f'export UV_CACHE_DIR=/tmp/uv_cache_$(id -u) && '
+            f'cd {self.sky_dir} && export HOME=$(pwd) && {cmd}')
+
+        cmd = (f'srun --unbuffered --quiet --overlap --jobid={self.job_id} '
+               f'--nodelist={self.slurm_node} '
+               f'--nodes=1 --ntasks=1 bash -c {shlex.quote(cmd)}')
+
+        return super().run(cmd, **kwargs)