skypilot-nightly 1.0.0.dev20251210__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
 
@@ -90,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:
@@ -133,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.
     #
@@ -148,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 = [
@@ -174,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()
@@ -233,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):
@@ -263,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
@@ -424,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.
 
@@ -442,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
@@ -622,9 +653,11 @@ class SSHCommandRunner(CommandRunner):
         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.
 
@@ -644,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
@@ -663,6 +698,7 @@ 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())
@@ -712,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,
@@ -763,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,
@@ -770,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.
 
@@ -810,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.
 
@@ -834,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))
@@ -872,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(
@@ -920,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))
@@ -1033,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.
@@ -1057,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
@@ -1093,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
@@ -1211,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)
@@ -1332,29 +1526,6 @@ class SlurmCommandRunner(SSHCommandRunner):
         self.job_id = job_id
         self.slurm_node = slurm_node
 
-        # Build a chained ProxyCommand that goes through the login node to reach
-        # the compute node where the job is running.
-
-        # First, build SSH options to reach the login node, using the user's
-        # existing proxy command if provided.
-        proxy_ssh_options = ' '.join(
-            ssh_options_list(self.ssh_private_key,
-                             None,
-                             ssh_proxy_command=self._ssh_proxy_command,
-                             port=self.port,
-                             disable_control_master=True))
-        login_node_proxy_command = (f'ssh {proxy_ssh_options} '
-                                    f'-W %h:%p {self.ssh_user}@{self.ip}')
-
-        # Update the proxy command to be the login node proxy, which will
-        # be used by super().run() to reach the compute node.
-        self._ssh_proxy_command = login_node_proxy_command
-        # Update self.ip to target the compute node.
-        self.ip = slurm_node
-        # Assume the compute node's SSH port is 22.
-        # TODO(kevin): Make this configurable if needed.
-        self.port = 22
-
     def rsync(
         self,
         source: str,
@@ -1365,40 +1536,35 @@ class SlurmCommandRunner(SSHCommandRunner):
         stream_logs: bool = True,
         max_retry: int = 1,
     ) -> None:
-        """Rsyncs files directly to the Slurm compute node,
-        by proxying through the Slurm login node.
-
-        For Slurm, files need to be accessible by compute nodes where jobs
-        execute via srun. This means either it has to be on the compute node's
-        local filesystem, or on a shared filesystem.
+        """Rsyncs files to/from the Slurm compute node using srun as transport.
         """
-        # TODO(kevin): We can probably optimize this to skip the proxying
-        # if the target dir is in a shared filesystem, since it will
-        # be accessible by the compute node.
-
-        # Build SSH options for rsync using the ProxyCommand set up in __init__
-        # to reach the compute node through the login node.
-        ssh_options = ' '.join(
-            ssh_options_list(
-                # Assume nothing and rely on default SSH behavior when -i is
-                # not specified.
-                None,
-                None,
-                ssh_proxy_command=self._ssh_proxy_command,
-                disable_control_master=True))
-        rsh_option = f'ssh {ssh_options}'
-
-        self._rsync(
-            source,
-            target,
-            # Compute node
-            node_destination=f'{self.ssh_user}@{self.slurm_node}',
-            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)
+        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
@@ -1420,14 +1586,6 @@ class SlurmCommandRunner(SSHCommandRunner):
         # could be part of a shared filesystem.
         # And similarly for SKY_RUNTIME_DIR. See constants.\
         # SKY_RUNTIME_DIR_ENV_VAR_KEY for more details.
-        #
-        # SSH directly to the compute node instead of using srun.
-        # This avoids Slurm's proctrack/cgroup which kills all processes
-        # when the job step ends (including child processes launched as
-        # a separate process group), breaking background process spawning
-        # (e.g., JobScheduler._run_job which uses launch_new_process_tree).
-        # Note: proctrack/cgroup is enabled by default on Nebius'
-        # Managed Soperator.
         cmd = (
             f'export {constants.SKY_RUNTIME_DIR_ENV_VAR_KEY}='
             f'"{self.skypilot_runtime_dir}" && '
@@ -1438,4 +1596,8 @@ class SlurmCommandRunner(SSHCommandRunner):
             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)