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

sky/utils/command_runner.py

@@ -5,20 +5,23 @@ import os
 import pathlib
 import shlex
 import time
-from typing import Any, Iterable, List, Optional, Tuple, Type, Union
+from typing import Any, Callable, Iterable, List, Optional, Tuple, Type, Union
 
 from sky import sky_logging
 from sky.skylet import constants
 from sky.skylet import log_lib
 from sky.utils import common_utils
+from sky.utils import control_master_utils
 from sky.utils import subprocess_utils
 from sky.utils import timeline
 
 logger = sky_logging.init_logger(__name__)
 
-# The git exclude file to support.
-GIT_EXCLUDE = '.git/info/exclude'
 # Rsync options
+# TODO(zhwu): This will print a per-file progress bar (with -P),
+# shooting a lot of messages to the output. --info=progress2 is used
+# to get a total progress bar, but it requires rsync>=3.1.0 and Mac
+# OS has a default rsync==2.6.9 (16 years old).
 RSYNC_DISPLAY_OPTION = '-Pavz'
 # Legend
 #   dir-merge: ignore file can appear in any subdir, applies to that
@@ -26,10 +29,14 @@ RSYNC_DISPLAY_OPTION = '-Pavz'
 # Note that "-" is mandatory for rsync and means all patterns in the ignore
 # files are treated as *exclude* patterns.  Non-exclude patterns, e.g., "!
 # do_not_exclude" doesn't work, even though git allows it.
-RSYNC_FILTER_OPTION = '--filter=\'dir-merge,- .gitignore\''
+RSYNC_FILTER_SKYIGNORE = f'--filter=\'dir-merge,- {constants.SKY_IGNORE_FILE}\''
+RSYNC_FILTER_GITIGNORE = f'--filter=\'dir-merge,- {constants.GIT_IGNORE_FILE}\''
+# The git exclude file to support.
+GIT_EXCLUDE = '.git/info/exclude'
 RSYNC_EXCLUDE_OPTION = '--exclude-from={}'
 
 _HASH_MAX_LENGTH = 10
+_DEFAULT_CONNECT_TIMEOUT = 30
 
 
 def _ssh_control_path(ssh_control_filename: Optional[str]) -> Optional[str]:
@@ -60,9 +67,12 @@ def ssh_options_list(
 ) -> List[str]:
     """Returns a list of sane options for 'ssh'."""
     if connect_timeout is None:
-        connect_timeout = 30
+        connect_timeout = _DEFAULT_CONNECT_TIMEOUT
     # Forked from Ray SSHOptions:
     # https://github.com/ray-project/ray/blob/master/python/ray/autoscaler/_private/command_runner.py
+    # Do not allow agent forwarding because SkyPilot API server has access to
+    # all user cluster private keys, which should not be all forwarded to
+    # individual user clusters.
     arg_dict = {
         # SSH port
         'Port': port,
@@ -75,7 +85,7 @@ def ssh_options_list(
         # that case.
         'UserKnownHostsFile': os.devnull,
         # Suppresses the warning messages, such as:
-        #   Warning: Permanently added '34.69.216.203' (ED25519) to the list of
+        #   Warning: Permanently added 'xx.xx.xx.xx' (EDxxx) to the list of
         #   known hosts.
         'LogLevel': 'ERROR',
         # Try fewer extraneous key pairs.
@@ -89,18 +99,25 @@ def ssh_options_list(
         'ServerAliveCountMax': 3,
         # ConnectTimeout.
         'ConnectTimeout': f'{connect_timeout}s',
-        # Agent forwarding for git.
-        'ForwardAgent': 'yes',
     }
     # 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
+    # ServerAliveInterval set, since the keepalive message may not be recognized
+    # by the custom proxy command, such as AWS SSM Session Manager.
+    #
     # We also do not use ControlMaster when we use `kubectl port-forward`
     # to access Kubernetes pods over SSH+Proxycommand. This is because the
     # process running ProxyCommand is kept running as long as the ssh session
     # is running and the ControlMaster keeps the session, which results in
     # 'ControlPersist' number of seconds delay per ssh commands ran.
     if (ssh_control_name is not None and docker_ssh_proxy_command is None and
-            not disable_control_master):
+            ssh_proxy_command is None and not disable_control_master):
         arg_dict.update({
             # Control path: important optimization as we do multiple ssh in one
             # sky.launch().
@@ -161,7 +178,7 @@ class CommandRunner:
         cmd: Union[str, List[str]],
         process_stream: bool,
         separate_stderr: bool,
-        skip_lines: int,
+        skip_num_lines: int,
         source_bashrc: bool = False,
     ) -> str:
         """Returns the command to run."""
@@ -170,7 +187,7 @@ class CommandRunner:
 
         # We need this to correctly run the cmd, and get the output.
         command = [
-            'bash',
+            '/bin/bash',
             '--login',
             '-c',
         ]
@@ -193,12 +210,12 @@ class CommandRunner:
             ]
         if not separate_stderr:
             command.append('2>&1')
-        if not process_stream and skip_lines:
+        if not process_stream and skip_num_lines:
             command += [
                 # A hack to remove the following bash warnings (twice):
                 #  bash: cannot set terminal process group
                 #  bash: no job control in this shell
-                f'| stdbuf -o0 tail -n +{skip_lines}',
+                f'| stdbuf -o0 tail -n +{skip_num_lines}',
                 # This is required to make sure the executor of command can get
                 # correct returncode, since linux pipe is used.
                 '; exit ${PIPESTATUS[0]}'
@@ -207,6 +224,111 @@ class CommandRunner:
         command_str = ' '.join(command)
         return command_str
 
+    def _rsync(
+            self,
+            source: str,
+            target: str,
+            node_destination: str,
+            up: bool,
+            rsh_option: str,
+            # Advanced options.
+            log_path: str = os.devnull,
+            stream_logs: bool = True,
+            max_retry: int = 1,
+            prefix_command: Optional[str] = None,
+            get_remote_home_dir: Callable[[], str] = lambda: '~') -> None:
+        """Builds the rsync command."""
+        # Build command.
+        rsync_command = []
+        if prefix_command is not None:
+            rsync_command.append(prefix_command)
+        rsync_command += ['rsync', RSYNC_DISPLAY_OPTION]
+
+        def _get_remote_home_dir_with_retry():
+            backoff = common_utils.Backoff(initial_backoff=1,
+                                           max_backoff_factor=5)
+            retries_left = max_retry
+            assert retries_left > 0, f'max_retry {max_retry} must be positive.'
+            while retries_left >= 0:
+                try:
+                    return get_remote_home_dir()
+                except Exception:  # pylint: disable=broad-except
+                    if retries_left == 0:
+                        raise
+                    sleep_time = backoff.current_backoff()
+                    logger.warning(f'Failed to get remote home dir '
+                                   f'- retrying in {sleep_time} seconds.')
+                    retries_left -= 1
+                    time.sleep(sleep_time)
+
+        # --filter
+        # The source is a local path, so we need to resolve it.
+        resolved_source = pathlib.Path(source).expanduser().resolve()
+        if (resolved_source / constants.SKY_IGNORE_FILE).exists():
+            rsync_command.append(RSYNC_FILTER_SKYIGNORE)
+        else:
+            rsync_command.append(RSYNC_FILTER_GITIGNORE)
+            if up:
+                # Build --exclude-from argument.
+                if (resolved_source / GIT_EXCLUDE).exists():
+                    # Ensure file exists; otherwise, rsync will error out.
+                    #
+                    # We shlex.quote() because the path may contain spaces:
+                    #   'my dir/.git/info/exclude'
+                    # Without quoting rsync fails.
+                    rsync_command.append(
+                        RSYNC_EXCLUDE_OPTION.format(
+                            shlex.quote(str(resolved_source / GIT_EXCLUDE))))
+
+        rsync_command.append(f'-e {shlex.quote(rsh_option)}')
+
+        if up:
+            resolved_target = target
+            if target.startswith('~'):
+                remote_home_dir = _get_remote_home_dir_with_retry()
+                resolved_target = target.replace('~', remote_home_dir)
+            full_source_str = str(resolved_source)
+            if resolved_source.is_dir():
+                full_source_str = os.path.join(full_source_str, '')
+            rsync_command.extend([
+                f'{full_source_str!r}',
+                f'{node_destination}:{resolved_target!r}',
+            ])
+        else:
+            resolved_source = source
+            if source.startswith('~'):
+                remote_home_dir = _get_remote_home_dir_with_retry()
+                resolved_source = source.replace('~', remote_home_dir)
+            rsync_command.extend([
+                f'{node_destination}:{resolved_source!r}',
+                f'{os.path.expanduser(target)!r}',
+            ])
+        command = ' '.join(rsync_command)
+        logger.debug(f'Running rsync command: {command}')
+
+        backoff = common_utils.Backoff(initial_backoff=5, max_backoff_factor=5)
+        assert max_retry > 0, f'max_retry {max_retry} must be positive.'
+        while max_retry >= 0:
+            returncode, stdout, stderr = log_lib.run_with_log(
+                command,
+                log_path=log_path,
+                stream_logs=stream_logs,
+                shell=True,
+                require_outputs=True)
+            if returncode == 0:
+                break
+            max_retry -= 1
+            time.sleep(backoff.current_backoff())
+
+        direction = 'up' if up else 'down'
+        error_msg = (f'Failed to rsync {direction}: {source} -> {target}. '
+                     'Ensure that the network is stable, then retry.')
+        subprocess_utils.handle_returncode(returncode,
+                                           command,
+                                           error_msg,
+                                           stderr=stdout + stderr,
+                                           stream_logs=stream_logs)
+
     @timeline.event
     def run(
             self,
@@ -222,7 +344,7 @@ class CommandRunner:
             separate_stderr: bool = False,
             connect_timeout: Optional[int] = None,
             source_bashrc: bool = False,
-            skip_lines: int = 0,
+            skip_num_lines: int = 0,
             **kwargs) -> Union[int, Tuple[int, str, str]]:
         """Runs the command on the cluster.
 
@@ -237,7 +359,7 @@ class CommandRunner:
             connect_timeout: timeout in seconds for the ssh connection.
             source_bashrc: Whether to source the ~/.bashrc before running the
                 command.
-            skip_lines: The number of lines to skip at the beginning of the
+            skip_num_lines: The number of lines to skip at the beginning of the
                 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.
@@ -293,6 +415,22 @@ class CommandRunner:
         returncode = self.run('true', connect_timeout=5, stream_logs=False)
         return returncode == 0
 
+    def close_cached_connection(self) -> None:
+        """Close the cached connection to the remote machine."""
+        pass
+
+    def port_forward_command(self,
+                             port_forward: List[Tuple[int, int]],
+                             connect_timeout: int = 1) -> List[str]:
+        """Command for forwarding ports from localhost to the remote machine.
+
+        Args:
+            port_forward: A list of ports to forward from the localhost to the
+                remote host.
+            connect_timeout: The timeout for the connection.
+        """
+        raise NotImplementedError
+
 
 class SSHCommandRunner(CommandRunner):
     """Runner for SSH commands."""
@@ -340,7 +478,9 @@ 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.disable_control_master = disable_control_master
+        self.disable_control_master = (
+            disable_control_master or
+            control_master_utils.should_disable_control_master())
         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}.')
@@ -359,9 +499,27 @@ class SSHCommandRunner(CommandRunner):
             self.port = port
             self._docker_ssh_proxy_command = None
 
-    def _ssh_base_command(self, *, ssh_mode: SshMode,
-                          port_forward: Optional[List[int]],
-                          connect_timeout: Optional[int]) -> List[str]:
+    def port_forward_command(self,
+                             port_forward: List[Tuple[int, int]],
+                             connect_timeout: int = 1) -> List[str]:
+        """Command for forwarding ports from localhost to the remote machine.
+
+        Args:
+            port_forward: A list of ports to forward from the local port to the
+                remote port.
+            connect_timeout: The timeout for the ssh connection.
+
+        Returns:
+            The command for forwarding ports from localhost to the remote
+            machine.
+        """
+        return self.ssh_base_command(ssh_mode=SshMode.INTERACTIVE,
+                                     port_forward=port_forward,
+                                     connect_timeout=connect_timeout)
+
+    def ssh_base_command(self, *, ssh_mode: SshMode,
+                         port_forward: Optional[List[Tuple[int, int]]],
+                         connect_timeout: Optional[int]) -> List[str]:
         ssh = ['ssh']
         if ssh_mode == SshMode.NON_INTERACTIVE:
             # Disable pseudo-terminal allocation. Otherwise, the output of
@@ -371,11 +529,10 @@ class SSHCommandRunner(CommandRunner):
             # Force pseudo-terminal allocation for interactive/login mode.
             ssh += ['-tt']
         if port_forward is not None:
-            for port in port_forward:
-                local = remote = port
+            for local, remote in port_forward:
                 logger.info(
                     f'Forwarding port {local} to port {remote} on localhost.')
-                ssh += ['-L', f'{remote}:localhost:{local}']
+                ssh += ['-NL', f'{remote}:localhost:{local}']
         if self._docker_ssh_proxy_command is not None:
             docker_ssh_proxy_command = self._docker_ssh_proxy_command(ssh)
         else:
@@ -391,13 +548,35 @@ class SSHCommandRunner(CommandRunner):
                 f'{self.ssh_user}@{self.ip}'
             ]
 
+    def close_cached_connection(self) -> None:
+        """Close the cached connection to the remote machine.
+
+        This is useful when we need to make the permission update effective of a
+        ssh user, e.g. usermod -aG docker $USER.
+        """
+        if self.ssh_control_name is not None:
+            control_path = _ssh_control_path(self.ssh_control_name)
+            if control_path is not None:
+                # Suppress the `Exit request sent.` output for this comamnd
+                # which would interrupt the CLI spinner.
+                cmd = (f'ssh -O exit -S {control_path}/%C '
+                       f'{self.ssh_user}@{self.ip} > /dev/null 2>&1')
+                logger.debug(f'Closing cached connection {control_path!r} with '
+                             f'cmd: {cmd}')
+                log_lib.run_with_log(cmd,
+                                     log_path=os.devnull,
+                                     require_outputs=False,
+                                     stream_logs=False,
+                                     process_stream=False,
+                                     shell=True)
+
     @timeline.event
     def run(
             self,
             cmd: Union[str, List[str]],
             *,
             require_outputs: bool = False,
-            port_forward: Optional[List[int]] = None,
+            port_forward: Optional[List[Tuple[int, int]]] = None,
             # Advanced options.
             log_path: str = os.devnull,
             # If False, do not redirect stdout/stderr to optimize performance.
@@ -407,7 +586,7 @@ class SSHCommandRunner(CommandRunner):
             separate_stderr: bool = False,
             connect_timeout: Optional[int] = None,
             source_bashrc: bool = False,
-            skip_lines: int = 0,
+            skip_num_lines: int = 0,
             **kwargs) -> Union[int, Tuple[int, str, str]]:
         """Uses 'ssh' to run 'cmd' on a node with ip.
 
@@ -428,7 +607,7 @@ class SSHCommandRunner(CommandRunner):
             connect_timeout: timeout in seconds for the ssh connection.
             source_bashrc: Whether to source the bashrc before running the
                 command.
-            skip_lines: The number of lines to skip at the beginning of the
+            skip_num_lines: The number of lines to skip at the beginning of the
                 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.
@@ -438,7 +617,7 @@ class SSHCommandRunner(CommandRunner):
             or
             A tuple of (returncode, stdout, stderr).
         """
-        base_ssh_command = self._ssh_base_command(
+        base_ssh_command = self.ssh_base_command(
             ssh_mode=ssh_mode,
             port_forward=port_forward,
             connect_timeout=connect_timeout)
@@ -451,7 +630,7 @@ class SSHCommandRunner(CommandRunner):
         command_str = self._get_command_to_run(cmd,
                                                process_stream,
                                                separate_stderr,
-                                               skip_lines=skip_lines,
+                                               skip_num_lines=skip_num_lines,
                                                source_bashrc=source_bashrc)
         command = base_ssh_command + [shlex.quote(command_str)]
 
@@ -506,30 +685,6 @@ class SSHCommandRunner(CommandRunner):
         Raises:
             exceptions.CommandError: rsync command failed.
         """
-        # Build command.
-        # TODO(zhwu): This will print a per-file progress bar (with -P),
-        # shooting a lot of messages to the output. --info=progress2 is used
-        # to get a total progress bar, but it requires rsync>=3.1.0 and Mac
-        # OS has a default rsync==2.6.9 (16 years old).
-        rsync_command = ['rsync', RSYNC_DISPLAY_OPTION]
-
-        # --filter
-        rsync_command.append(RSYNC_FILTER_OPTION)
-
-        if up:
-            # The source is a local path, so we need to resolve it.
-            # --exclude-from
-            resolved_source = pathlib.Path(source).expanduser().resolve()
-            if (resolved_source / GIT_EXCLUDE).exists():
-                # Ensure file exists; otherwise, rsync will error out.
-                #
-                # We shlex.quote() because the path may contain spaces:
-                #   'my dir/.git/info/exclude'
-                # Without quoting rsync fails.
-                rsync_command.append(
-                    RSYNC_EXCLUDE_OPTION.format(
-                        shlex.quote(str(resolved_source / GIT_EXCLUDE))))
-
         if self._docker_ssh_proxy_command is not None:
             docker_ssh_proxy_command = self._docker_ssh_proxy_command(['ssh'])
         else:
@@ -542,43 +697,251 @@ class SSHCommandRunner(CommandRunner):
                 docker_ssh_proxy_command=docker_ssh_proxy_command,
                 port=self.port,
                 disable_control_master=self.disable_control_master))
-        rsync_command.append(f'-e "ssh {ssh_options}"')
-        # To support spaces in the path, we need to quote source and target.
-        # rsync doesn't support '~' in a quoted local path, but it is ok to
-        # have '~' in a quoted remote path.
-        if up:
-            full_source_str = str(resolved_source)
-            if resolved_source.is_dir():
-                full_source_str = os.path.join(full_source_str, '')
-            rsync_command.extend([
-                f'{full_source_str!r}',
-                f'{self.ssh_user}@{self.ip}:{target!r}',
-            ])
-        else:
-            rsync_command.extend([
-                f'{self.ssh_user}@{self.ip}:{source!r}',
-                f'{os.path.expanduser(target)!r}',
-            ])
-        command = ' '.join(rsync_command)
+        rsh_option = f'ssh {ssh_options}'
+        self._rsync(source,
+                    target,
+                    node_destination=f'{self.ssh_user}@{self.ip}',
+                    up=up,
+                    rsh_option=rsh_option,
+                    log_path=log_path,
+                    stream_logs=stream_logs,
+                    max_retry=max_retry)
 
-        backoff = common_utils.Backoff(initial_backoff=5, max_backoff_factor=5)
-        while max_retry >= 0:
-            returncode, stdout, stderr = log_lib.run_with_log(
-                command,
-                log_path=log_path,
-                stream_logs=stream_logs,
-                shell=True,
-                require_outputs=True)
-            if returncode == 0:
-                break
-            max_retry -= 1
-            time.sleep(backoff.current_backoff())
 
-        direction = 'up' if up else 'down'
-        error_msg = (f'Failed to rsync {direction}: {source} -> {target}. '
-                     'Ensure that the network is stable, then retry.')
-        subprocess_utils.handle_returncode(returncode,
-                                           command,
-                                           error_msg,
-                                           stderr=stdout + stderr,
-                                           stream_logs=stream_logs)
+class KubernetesCommandRunner(CommandRunner):
+    """Runner for Kubernetes commands."""
+
+    _MAX_RETRIES_FOR_RSYNC = 3
+
+    def __init__(
+        self,
+        node: Tuple[Tuple[str, Optional[str]], str],
+        **kwargs,
+    ):
+        """Initialize KubernetesCommandRunner.
+
+        Example Usage:
+            runner = KubernetesCommandRunner((namespace, context), pod_name))
+            runner.run('ls -l')
+            runner.rsync(source, target, up=True)
+
+        Args:
+            node: The namespace and pod_name of the remote machine.
+        """
+        del kwargs
+        super().__init__(node)
+        (self.namespace, self.context), self.pod_name = node
+
+    @property
+    def node_id(self) -> str:
+        return f'{self.context}-{self.namespace}-{self.pod_name}'
+
+    def port_forward_command(self,
+                             port_forward: List[Tuple[int, int]],
+                             connect_timeout: int = 1) -> List[str]:
+        """Command for forwarding ports from localhost to the remote machine.
+
+        Args:
+            port_forward: A list of ports to forward from the local port to the
+                remote port. Currently, only one port is supported, i.e. the
+                list should have only one element.
+            connect_timeout: The timeout for the ssh connection.
+        """
+        assert port_forward and len(port_forward) == 1, (
+            'Only one port is supported for Kubernetes port-forward.')
+        kubectl_args = [
+            '--pod-running-timeout', f'{connect_timeout}s', '-n', self.namespace
+        ]
+        if self.context:
+            kubectl_args += ['--context', self.context]
+        local_port, remote_port = port_forward[0]
+        local_port_str = f'{local_port}' if local_port is not None else ''
+        kubectl_cmd = [
+            'kubectl',
+            *kubectl_args,
+            'port-forward',
+            f'pod/{self.pod_name}',
+            f'{local_port_str}:{remote_port}',
+        ]
+        return kubectl_cmd
+
+    @timeline.event
+    def run(
+            self,
+            cmd: Union[str, List[str]],
+            *,
+            port_forward: Optional[List[int]] = None,
+            require_outputs: bool = False,
+            # Advanced options.
+            log_path: str = os.devnull,
+            # If False, do not redirect stdout/stderr to optimize performance.
+            process_stream: bool = True,
+            stream_logs: bool = True,
+            ssh_mode: SshMode = SshMode.NON_INTERACTIVE,
+            separate_stderr: bool = False,
+            connect_timeout: Optional[int] = None,
+            source_bashrc: bool = False,
+            skip_num_lines: int = 0,
+            **kwargs) -> Union[int, Tuple[int, str, str]]:
+        """Uses 'kubectl exec' to run 'cmd' on a pod by its name and namespace.
+
+        Args:
+            cmd: The command to run.
+            port_forward: This should be None for k8s.
+
+            Advanced options:
+
+            require_outputs: Whether to return the stdout/stderr of the command.
+            log_path: Redirect stdout/stderr to the log_path.
+            stream_logs: Stream logs to the stdout/stderr.
+            check: Check the success of the command.
+            ssh_mode: The mode to use for ssh.
+                See SSHMode for more details.
+            separate_stderr: Whether to separate stderr from stdout.
+            connect_timeout: timeout in seconds for the pod connection.
+            source_bashrc: Whether to source the bashrc before running the
+                command.
+            skip_num_lines: The number of lines to skip at the beginning of the
+                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.
+
+
+        Returns:
+            returncode
+            or
+            A tuple of (returncode, stdout, stderr).
+        """
+        # TODO(zhwu): implement port_forward for k8s.
+        assert port_forward is None, ('port_forward is not supported for k8s '
+                                      f'for now, but got: {port_forward}')
+        if connect_timeout is None:
+            connect_timeout = _DEFAULT_CONNECT_TIMEOUT
+        kubectl_args = [
+            '--pod-running-timeout', f'{connect_timeout}s', '-n', self.namespace
+        ]
+        if self.context:
+            kubectl_args += ['--context', self.context]
+        # If context is none, it means we are using incluster auth. In this
+        # case, need to set KUBECONFIG to /dev/null to avoid using kubeconfig.
+        if self.context is None:
+            kubectl_args += ['--kubeconfig', '/dev/null']
+        kubectl_args += [self.pod_name]
+        if ssh_mode == SshMode.LOGIN:
+            assert isinstance(cmd, list), 'cmd must be a list for login mode.'
+            base_cmd = ['kubectl', 'exec', '-it', *kubectl_args, '--']
+            command = base_cmd + cmd
+            proc = subprocess_utils.run(command, shell=False, check=False)
+            return proc.returncode, '', ''
+
+        kubectl_base_command = ['kubectl', 'exec']
+
+        if ssh_mode == SshMode.INTERACTIVE:
+            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 = 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
+            # after '--' will not work for some commands, such as '&&', '>' etc.
+            '/bin/bash',
+            '-c',
+            shlex.quote(command_str)
+        ]
+
+        log_dir = os.path.expanduser(os.path.dirname(log_path))
+        os.makedirs(log_dir, exist_ok=True)
+
+        executable = None
+        if not process_stream:
+            if stream_logs:
+                command += [
+                    f'| tee {log_path}',
+                    # This also requires the executor to be '/bin/bash' instead
+                    # of the default '/bin/sh'.
+                    '; exit ${PIPESTATUS[0]}'
+                ]
+            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)
+
+    @timeline.event
+    def rsync(
+        self,
+        source: str,
+        target: str,
+        *,
+        up: bool,
+        # Advanced options.
+        log_path: str = os.devnull,
+        stream_logs: bool = True,
+        max_retry: int = _MAX_RETRIES_FOR_RSYNC,
+    ) -> None:
+        """Uses 'rsync' to sync 'source' to 'target'.
+
+        Args:
+            source: The source path.
+            target: The target path.
+            up: The direction of the sync, True for local to cluster, False
+              for cluster to local.
+            log_path: Redirect stdout/stderr to the log_path.
+            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.
+
+        Raises:
+            exceptions.CommandError: rsync command failed.
+        """
+
+        def get_remote_home_dir() -> str:
+            # Use `echo ~` to get the remote home directory, instead of pwd or
+            # echo $HOME, because pwd can be `/` when the remote user is root
+            # and $HOME is not always set.
+            rc, remote_home_dir, stderr = self.run('echo ~',
+                                                   require_outputs=True,
+                                                   separate_stderr=True,
+                                                   stream_logs=False)
+            if rc != 0:
+                raise ValueError('Failed to get remote home directory: '
+                                 f'{remote_home_dir + stderr}')
+            remote_home_dir = remote_home_dir.strip()
+            return remote_home_dir
+
+        # Build command.
+        helper_path = os.path.join(os.path.abspath(os.path.dirname(__file__)),
+                                   'kubernetes', 'rsync_helper.sh')
+        namespace_context = f'{self.namespace}+{self.context}'
+        # Avoid rsync interpreting :, /, and + in namespace_context as the
+        # default delimiter for options and arguments.
+        # rsync_helper.sh will parse the namespace_context by reverting the
+        # encoding and pass it to kubectl exec.
+        encoded_namespace_context = (namespace_context.replace(
+            '@', '%40').replace(':', '%3A').replace('/',
+                                                    '%2F').replace('+', '%2B'))
+        self._rsync(
+            source,
+            target,
+            node_destination=f'{self.pod_name}@{encoded_namespace_context}',
+            up=up,
+            rsh_option=helper_path,
+            log_path=log_path,
+            stream_logs=stream_logs,
+            max_retry=max_retry,
+            prefix_command=f'chmod +x {helper_path} && ',
+            # rsync with `kubectl` as the rsh command will cause ~/xx parsed as
+            # /~/xx, so we need to replace ~ with the remote home directory. We
+            # only need to do this when ~ is at the beginning of the path.
+            get_remote_home_dir=get_remote_home_dir)