skypilot-nightly 1.0.0.dev20250910__py3-none-any.whl → 1.0.0.dev20250913__py3-none-any.whl

sky/utils/command_runner.py

@@ -469,15 +469,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 +596,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 +623,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
@@ -646,22 +655,28 @@ class SSHCommandRunner(CommandRunner):
             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 +695,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 +913,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,7 +925,10 @@ 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 = [

sky/utils/command_runner.pyi

@@ -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,
@@ -200,6 +208,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 +287,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/context_utils.py

@@ -10,6 +10,8 @@ import sys
 import typing
 from typing import Any, Callable, IO, Optional, Tuple, TypeVar
 
+from typing_extensions import ParamSpec
+
 from sky import sky_logging
 from sky.utils import context
 from sky.utils import subprocess_utils
@@ -173,9 +175,14 @@ def cancellation_guard(func: F) -> F:
     return typing.cast(F, wrapper)
 
 
+P = ParamSpec('P')
+T = TypeVar('T')
+
+
 # TODO(aylei): replace this with asyncio.to_thread once we drop support for
 # python 3.8
-def to_thread(func, /, *args, **kwargs):
+def to_thread(func: Callable[P, T], /, *args: P.args,
+              **kwargs: P.kwargs) -> 'asyncio.Future[T]':
     """Asynchronously run function *func* in a separate thread.
 
     This is same as asyncio.to_thread added in python 3.9
@@ -183,5 +190,11 @@ def to_thread(func, /, *args, **kwargs):
     loop = asyncio.get_running_loop()
     # This is critical to pass the current coroutine context to the new thread
     pyctx = contextvars.copy_context()
-    func_call = functools.partial(pyctx.run, func, *args, **kwargs)
+    func_call: Callable[..., T] = functools.partial(
+        # partial deletes arguments type and thus can't figure out the return
+        # type of pyctx.run
+        pyctx.run,  # type: ignore
+        func,
+        *args,
+        **kwargs)
     return loop.run_in_executor(None, func_call)

sky/utils/controller_utils.py

@@ -228,15 +228,21 @@ def get_controller_for_pool(pool: bool) -> Controllers:
 def high_availability_specified(cluster_name: Optional[str]) -> bool:
     """Check if the controller high availability is specified in user config.
     """
-    # pylint: disable=import-outside-toplevel
-    from sky.jobs import utils as managed_job_utils
-    if managed_job_utils.is_consolidation_mode():
-        return True
-
     controller = Controllers.from_name(cluster_name)
     if controller is None:
         return False
 
+    if controller.value.controller_type == 'jobs':
+        # pylint: disable-next=import-outside-toplevel
+        from sky.jobs import utils as managed_job_utils
+        if managed_job_utils.is_consolidation_mode():
+            return True
+    elif controller.value.controller_type == 'serve':
+        # pylint: disable-next=import-outside-toplevel
+        from sky.serve import serve_utils
+        if serve_utils.is_consolidation_mode():
+            return True
+
     if skypilot_config.loaded():
         return skypilot_config.get_nested((controller.value.controller_type,
                                            'controller', 'high_availability'),

sky/utils/db/migration_utils.py

@@ -17,7 +17,7 @@ logger = sky_logging.init_logger(__name__)
 DB_INIT_LOCK_TIMEOUT_SECONDS = 10
 
 GLOBAL_USER_STATE_DB_NAME = 'state_db'
-GLOBAL_USER_STATE_VERSION = '007'
+GLOBAL_USER_STATE_VERSION = '008'
 GLOBAL_USER_STATE_LOCK_PATH = '~/.sky/locks/.state_db.lock'
 
 SPOT_JOBS_DB_NAME = 'spot_jobs_db'