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

sky/utils/cli_utils/status_utils.py

@@ -11,6 +11,7 @@ from sky.utils import common_utils
 from sky.utils import log_utils
 from sky.utils import resources_utils
 from sky.utils import status_lib
+from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
     from sky.provision.kubernetes import utils as kubernetes_utils
@@ -105,11 +106,9 @@ def show_status_table(cluster_records: List[responses.StatusResponse],
 
     if query_clusters:
         cluster_names = {record['name'] for record in cluster_records}
-        not_found_clusters = [
-            repr(cluster)
-            for cluster in query_clusters
-            if cluster not in cluster_names
-        ]
+        not_found_clusters = ux_utils.get_non_matched_query(
+            query_clusters, cluster_names)
+        not_found_clusters = [repr(cluster) for cluster in not_found_clusters]
         if not_found_clusters:
             cluster_str = 'Cluster'
             if len(not_found_clusters) > 1:
@@ -283,8 +282,14 @@ def _get_resources(cluster_record: _ClusterRecord,
             if resources_str_full is not None:
                 resources_str = resources_str_full
         if resources_str is None:
-            resources_str = resources_utils.get_readable_resources_repr(
-                handle, simplify=truncate)
+            resources_str_simple, resources_str_full = (
+                resources_utils.get_readable_resources_repr(
+                    handle, simplified_only=truncate))
+            if truncate:
+                resources_str = resources_str_simple
+            else:
+                assert resources_str_full is not None
+                resources_str = resources_str_full
 
         return resources_str
     return '-'

sky/utils/cluster_utils.py

@@ -144,6 +144,9 @@ class SSHConfigHelper(object):
             username = docker_user
 
         key_path = cls.generate_local_key_file(cluster_name, auth_config)
+        # Keep the unexpanded path for SSH config (with ~)
+        key_path_for_config = key_path
+        # Expand the path for internal operations that need absolute path
         key_path = os.path.expanduser(key_path)
         sky_autogen_comment = ('# Added by sky (use `sky stop/down '
                                f'{cluster_name}` to remove)')
@@ -190,11 +193,29 @@ class SSHConfigHelper(object):
         proxy_command = auth_config.get('ssh_proxy_command', None)
 
         docker_proxy_command_generator = None
+        proxy_command_for_nodes = proxy_command
         if docker_user is not None:
-            docker_proxy_command_generator = lambda ip, port: ' '.join(
-                ['ssh'] + command_runner.ssh_options_list(
-                    key_path, ssh_control_name=None, port=port) +
-                ['-W', '%h:%p', f'{auth_config["ssh_user"]}@{ip}'])
+
+            def _docker_proxy_cmd(ip: str, port: int) -> str:
+                inner_proxy = proxy_command
+                inner_port = port or 22
+                if inner_proxy is not None:
+                    inner_proxy = inner_proxy.replace('%h', ip)
+                    inner_proxy = inner_proxy.replace('%p', str(inner_port))
+                return ' '.join(['ssh'] + command_runner.ssh_options_list(
+                    key_path,
+                    ssh_control_name=None,
+                    ssh_proxy_command=inner_proxy,
+                    port=inner_port,
+                    # ProxyCommand (ssh -W) is a forwarding tunnel, not an
+                    # interactive session. ControlMaster would cache these
+                    # processes, causing them to hang and block subsequent
+                    # connections. Each ProxyCommand should be ephemeral.
+                    disable_control_master=True
+                ) + ['-W', '%h:%p', f'{auth_config["ssh_user"]}@{ip}'])
+
+            docker_proxy_command_generator = _docker_proxy_cmd
+            proxy_command_for_nodes = None
 
         codegen = ''
         # Add the nodes to the codegen
@@ -208,8 +229,9 @@ class SSHConfigHelper(object):
             node_name = cluster_name if i == 0 else cluster_name + f'-worker{i}'
             # TODO(romilb): Update port number when k8s supports multinode
             codegen += cls._get_generated_config(
-                sky_autogen_comment, node_name, ip, username, key_path,
-                proxy_command, port, docker_proxy_command) + '\n'
+                sky_autogen_comment, node_name, ip, username,
+                key_path_for_config, proxy_command_for_nodes, port,
+                docker_proxy_command) + '\n'
 
         cluster_config_path = os.path.expanduser(
             cls.ssh_cluster_path.format(cluster_name))

sky/utils/command_runner.py

@@ -3,6 +3,7 @@ import enum
 import hashlib
 import os
 import pathlib
+import re
 import shlex
 import sys
 import time
@@ -13,6 +14,7 @@ from sky import exceptions
 from sky import sky_logging
 from sky.skylet import constants
 from sky.skylet import log_lib
+from sky.utils import auth_utils
 from sky.utils import common_utils
 from sky.utils import context_utils
 from sky.utils import control_master_utils
@@ -22,6 +24,9 @@ from sky.utils import timeline
 
 logger = sky_logging.init_logger(__name__)
 
+# Pattern to extract home directory from command output
+_HOME_DIR_PATTERN = re.compile(r'SKYPILOT_HOME_DIR: ([^\s\n]+)')
+
 # 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
@@ -183,17 +188,25 @@ class CommandRunner:
         return '-'.join(str(x) for x in self.node)
 
     def _get_remote_home_dir(self) -> 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)
+        # Use pattern matching to extract home directory.
+        # Some container images print MOTD when login shells start, which can
+        # contaminate command output. We use a unique pattern to extract the
+        # actual home directory reliably.
+        rc, output, stderr = self.run('echo "SKYPILOT_HOME_DIR: $(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()
+                             f'{output + stderr}')
+
+        # Extract home directory using pattern matching
+        home_dir_match = _HOME_DIR_PATTERN.search(output)
+        if home_dir_match:
+            remote_home_dir = home_dir_match.group(1)
+        else:
+            raise ValueError('Failed to find remote home directory identifier: '
+                             f'{output + stderr}')
         return remote_home_dir
 
     def _get_command_to_run(
@@ -414,7 +427,6 @@ class CommandRunner:
                 SkyPilot but we still want to get rid of some warning messages,
                 such as SSH warnings.
 
-
         Returns:
             returncode
             or
@@ -469,15 +481,19 @@ class CommandRunner:
         """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]:
+    def port_forward_command(
+            self,
+            port_forward: List[Tuple[int, int]],
+            connect_timeout: int = 1,
+            ssh_mode: SshMode = SshMode.INTERACTIVE) -> 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.
+            ssh_mode: The mode to use for ssh.
+                See SSHMode for more details.
         """
         raise NotImplementedError
 
@@ -592,6 +608,7 @@ class SSHCommandRunner(CommandRunner):
         ssh_proxy_command: Optional[str] = None,
         docker_user: Optional[str] = None,
         disable_control_master: Optional[bool] = False,
+        port_forward_execute_remote_command: Optional[bool] = False,
     ):
         """Initialize SSHCommandRunner.
 
@@ -618,6 +635,10 @@ class SSHCommandRunner(CommandRunner):
             disable_control_master: bool; specifies either or not the ssh
                 command will utilize ControlMaster. We currently disable
                 it for k8s instance.
+            port_forward_execute_remote_command: bool; specifies whether to
+                add -N to the port forwarding command. This is useful if you
+                want to run a command on the remote machine to make sure the
+                SSH tunnel is established.
         """
         super().__init__(node)
         ip, port = node
@@ -629,39 +650,63 @@ class SSHCommandRunner(CommandRunner):
         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)
         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}.')
-            # Already checked in resources
-            assert ssh_proxy_command is None, (
-                'ssh_proxy_command is not supported when using docker.')
+            # When connecting via docker, the outer SSH hop points to the
+            # container's sshd (localhost). Preserve the user proxy for the
+            # inner hop that reaches the host VM, and clear the outer proxy to
+            # avoid forwarding localhost through the jump host.
+            inner_proxy_command = ssh_proxy_command
+            inner_proxy_port = port or 22
+            self._ssh_proxy_command = None
             self.ip = 'localhost'
             self.ssh_user = docker_user
             self.port = constants.DEFAULT_DOCKER_PORT
+            if inner_proxy_command is not None:
+                # Replace %h/%p placeholders with actual host values, since the
+                # final destination from the perspective of the user proxy is
+                # the host VM (ip, inner_proxy_port).
+                inner_proxy_command = inner_proxy_command.replace('%h', ip)
+                inner_proxy_command = inner_proxy_command.replace(
+                    '%p', str(inner_proxy_port))
             self._docker_ssh_proxy_command = lambda ssh: ' '.join(
-                ssh + ssh_options_list(ssh_private_key, None
-                                      ) + ['-W', '%h:%p', f'{ssh_user}@{ip}'])
+                ssh + ssh_options_list(ssh_private_key,
+                                       None,
+                                       ssh_proxy_command=inner_proxy_command,
+                                       port=inner_proxy_port,
+                                       disable_control_master=self.
+                                       disable_control_master) +
+                ['-W', '%h:%p', f'{ssh_user}@{ip}'])
         else:
             self.ip = ip
             self.ssh_user = ssh_user
             self.port = port
             self._docker_ssh_proxy_command = None
+        self.port_forward_execute_remote_command = (
+            port_forward_execute_remote_command)
 
-    def port_forward_command(self,
-                             port_forward: List[Tuple[int, int]],
-                             connect_timeout: int = 1) -> List[str]:
+    def port_forward_command(
+            self,
+            port_forward: List[Tuple[int, int]],
+            connect_timeout: int = 1,
+            ssh_mode: SshMode = SshMode.INTERACTIVE) -> 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.
+            ssh_mode: The mode to use for ssh.
+                See SSHMode for more details.
 
         Returns:
             The command for forwarding ports from localhost to the remote
             machine.
         """
-        return self.ssh_base_command(ssh_mode=SshMode.INTERACTIVE,
+        return self.ssh_base_command(ssh_mode=ssh_mode,
                                      port_forward=port_forward,
                                      connect_timeout=connect_timeout)
 
@@ -680,7 +725,11 @@ class SSHCommandRunner(CommandRunner):
             for local, remote in port_forward:
                 logger.debug(
                     f'Forwarding local port {local} to remote port {remote}.')
-                ssh += ['-NL', f'{local}:localhost:{remote}']
+                if self.port_forward_execute_remote_command:
+                    ssh += ['-L']
+                else:
+                    ssh += ['-NL']
+                ssh += [f'{local}:localhost:{remote}']
         if self._docker_ssh_proxy_command is not None:
             docker_ssh_proxy_command = self._docker_ssh_proxy_command(ssh)
         else:
@@ -894,9 +943,11 @@ class KubernetesCommandRunner(CommandRunner):
         else:
             return f'pod/{self.pod_name}'
 
-    def port_forward_command(self,
-                             port_forward: List[Tuple[int, int]],
-                             connect_timeout: int = 1) -> List[str]:
+    def port_forward_command(
+            self,
+            port_forward: List[Tuple[int, int]],
+            connect_timeout: int = 1,
+            ssh_mode: SshMode = SshMode.INTERACTIVE) -> List[str]:
         """Command for forwarding ports from localhost to the remote machine.
 
         Args:
@@ -904,14 +955,25 @@ class KubernetesCommandRunner(CommandRunner):
                 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.
+            ssh_mode: The mode to use for ssh.
+                See SSHMode for more details.
         """
+        del ssh_mode  # unused
         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
         ]
+        # The same logic to either set `--context` to the k8s context where
+        # the sky cluster is hosted, or `--kubeconfig` to /dev/null for
+        # in-cluster k8s is used below in the `run()` method.
         if self.context:
             kubectl_args += ['--context', self.context]
+        # If context is none, it means the cluster is hosted on in-cluster k8s.
+        # In this case, we need to set KUBECONFIG to /dev/null to avoid looking
+        # for the cluster in whatever active context is set in the kubeconfig.
+        else:
+            kubectl_args += ['--kubeconfig', '/dev/null']
         local_port, remote_port = port_forward[0]
         local_port_str = f'{local_port}' if local_port is not None else ''
 
@@ -967,7 +1029,6 @@ class KubernetesCommandRunner(CommandRunner):
                 SkyPilot but we still want to get rid of some warning messages,
                 such as SSH warnings.
 
-
         Returns:
             returncode
             or

sky/utils/command_runner.pyi

@@ -36,9 +36,9 @@ def ssh_options_list(
 
 
 class SshMode(enum.Enum):
-    NON_INTERACTIVE: int
-    INTERACTIVE: int
-    LOGIN: int
+    NON_INTERACTIVE = ...
+    INTERACTIVE = ...
+    LOGIN = ...
 
 
 class CommandRunner:
@@ -106,6 +106,13 @@ class CommandRunner:
               max_retry: int = ...) -> None:
         ...
 
+    def port_forward_command(
+            self,
+            port_forward: List[Tuple[int, int]],
+            connect_timeout: int = 1,
+            ssh_mode: SshMode = SshMode.INTERACTIVE) -> List[str]:
+        ...
+
     @classmethod
     def make_runner_list(cls: typing.Type[CommandRunner],
                          node_list: Iterable[Tuple[Any, ...]],
@@ -127,6 +134,7 @@ class SSHCommandRunner(CommandRunner):
     ssh_control_name: Optional[str]
     docker_user: str
     disable_control_master: Optional[bool]
+    port_forward_execute_remote_command: Optional[bool]
 
     def __init__(
         self,
@@ -134,8 +142,10 @@ class SSHCommandRunner(CommandRunner):
         ssh_user: str,
         ssh_private_key: str,
         ssh_control_name: Optional[str] = ...,
+        ssh_proxy_command: Optional[str] = ...,
         docker_user: Optional[str] = ...,
         disable_control_master: Optional[bool] = ...,
+        port_forward_execute_remote_command: Optional[bool] = ...,
     ) -> None:
         ...
 
@@ -190,6 +200,15 @@ class SSHCommandRunner(CommandRunner):
             **kwargs) -> Union[Tuple[int, str, str], int]:
         ...
 
+    def ssh_base_command(
+        self,
+        *,
+        ssh_mode: SshMode,
+        port_forward: Optional[List[Tuple[int, int]]],
+        connect_timeout: Optional[int],
+    ) -> List[str]:
+        ...
+
     def rsync(self,
               source: str,
               target: str,
@@ -200,6 +219,13 @@ class SSHCommandRunner(CommandRunner):
               max_retry: int = ...) -> None:
         ...
 
+    def port_forward_command(
+            self,
+            port_forward: List[Tuple[int, int]],
+            connect_timeout: int = 1,
+            ssh_mode: SshMode = SshMode.INTERACTIVE) -> List[str]:
+        ...
+
 
 class KubernetesCommandRunner(CommandRunner):
 
@@ -272,6 +298,13 @@ class KubernetesCommandRunner(CommandRunner):
               max_retry: int = ...) -> None:
         ...
 
+    def port_forward_command(
+            self,
+            port_forward: List[Tuple[int, int]],
+            connect_timeout: int = 1,
+            ssh_mode: SshMode = SshMode.INTERACTIVE) -> List[str]:
+        ...
+
 
 class LocalProcessCommandRunner(CommandRunner):
 

sky/utils/common.py

@@ -31,7 +31,7 @@ JOB_CONTROLLER_NAME: str
 def refresh_server_id() -> None:
     """Refresh the server id.
 
-    This function is used to ensure the server id is read from the authorative
+    This function is used to ensure the server id is read from the authoritative
     source.
     """
     global SERVER_ID
@@ -42,6 +42,8 @@ def refresh_server_id() -> None:
     JOB_CONTROLLER_NAME = f'{JOB_CONTROLLER_PREFIX}{SERVER_ID}'
 
 
+# TODO(kevin): Remove this side effect and have callers call
+# refresh_server_id() explicitly as needed.
 refresh_server_id()
 
 

sky/utils/common_utils.py

@@ -1,8 +1,10 @@
 """Utils shared between all of sky"""
 
+import ctypes
 import difflib
 import enum
 import functools
+import gc
 import getpass
 import hashlib
 import inspect
@@ -263,13 +265,16 @@ def get_global_job_id(job_timestamp: str,
 
 class Backoff:
     """Exponential backoff with jittering."""
-    MULTIPLIER = 1.6
     JITTER = 0.4
 
-    def __init__(self, initial_backoff: float = 5, max_backoff_factor: int = 5):
+    def __init__(self,
+                 initial_backoff: float = 5,
+                 max_backoff_factor: int = 5,
+                 multiplier: float = 1.6):
         self._initial = True
         self._backoff = 0.0
         self._initial_backoff = initial_backoff
+        self._multiplier = multiplier
         self._max_backoff = max_backoff_factor * self._initial_backoff
 
     # https://github.com/grpc/grpc/blob/2d4f3c56001cd1e1f85734b2f7c5ce5f2797c38a/doc/connection-backoff.md
@@ -281,7 +286,7 @@ class Backoff:
             self._initial = False
             self._backoff = min(self._initial_backoff, self._max_backoff)
         else:
-            self._backoff = min(self._backoff * self.MULTIPLIER,
+            self._backoff = min(self._backoff * self._multiplier,
                                 self._max_backoff)
         self._backoff += random.uniform(-self.JITTER * self._backoff,
                                         self.JITTER * self._backoff)
@@ -994,7 +999,17 @@ def get_mem_size_gb() -> float:
         except ValueError as e:
             with ux_utils.print_exception_no_traceback():
                 raise ValueError(
-                    f'Failed to parse the memory size from {mem_size}') from e
+                    f'Failed to parse the memory size from {mem_size} (GB)'
+                ) from e
+    mem_size = os.getenv('SKYPILOT_POD_MEMORY_BYTES_LIMIT')
+    if mem_size is not None:
+        try:
+            return float(mem_size) / (1024**3)
+        except ValueError as e:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    f'Failed to parse the memory size from {mem_size} (bytes)'
+                ) from e
     return _mem_size_gb()
 
 
@@ -1090,3 +1105,21 @@ def removeprefix(string: str, prefix: str) -> str:
     if string.startswith(prefix):
         return string[len(prefix):]
     return string
+
+
+def release_memory():
+    """Release the process memory"""
+    # Do the best effort to release the python heap and let malloc_trim
+    # be more efficient.
+    try:
+        gc.collect()
+        if sys.platform.startswith('linux'):
+            # Will fail on musl (alpine), but at least it works on our
+            # official docker images.
+            libc = ctypes.CDLL('libc.so.6')
+            return libc.malloc_trim(0)
+        return 0
+    except Exception as e:  # pylint: disable=broad-except
+        logger.error(f'Failed to release memory: '
+                     f'{format_exception(e)}')
+        return 0

sky/utils/config_utils.py

@@ -272,7 +272,7 @@ def get_cloud_config_value_from_dict(
     """
     input_config = Config(dict_config)
     region_key = None
-    if cloud == 'kubernetes':
+    if cloud in ('kubernetes', 'ssh'):
         region_key = 'context_configs'
     elif cloud in _REGION_CONFIG_CLOUDS:
         region_key = 'region_configs'
@@ -283,19 +283,6 @@ def get_cloud_config_value_from_dict(
             keys=(cloud, region_key, region) + keys,
             default_value=None,
             override_configs=override_configs)
-        if not per_context_config and cloud in _REGION_CONFIG_CLOUDS:
-            # TODO (kyuds): Backward compatibility, remove after 0.11.0.
-            per_context_config = input_config.get_nested(
-                keys=(cloud, region) + keys,
-                default_value=None,
-                override_configs=override_configs)
-            if per_context_config is not None:
-                logger.info(
-                    f'{cloud} configuration is using the legacy format. \n'
-                    'This format will be deprecated after 0.11.0, refer to '
-                    '`https://docs.skypilot.co/en/latest/reference/config.html` '  # pylint: disable=line-too-long
-                    'for the new format. Please use `region_configs` to specify region specific configuration.'
-                )
     # if no override found for specified region
     general_config = input_config.get_nested(keys=(cloud,) + keys,
                                              default_value=default_value,