skypilot-nightly 1.0.0.dev20250509__py3-none-any.whl → 1.0.0.dev20251107__py3-none-any.whl

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):
 
@@ -271,3 +297,67 @@ class KubernetesCommandRunner(CommandRunner):
               stream_logs: bool = ...,
               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):
+
+    def __init__(self) -> None:
+        ...
+
+    @typing.overload
+    def run(self,
+            cmd: Union[str, List[str]],
+            *,
+            port_forward: Optional[List[int]] = ...,
+            require_outputs: Literal[False] = ...,
+            log_path: str = ...,
+            process_stream: bool = ...,
+            stream_logs: bool = ...,
+            ssh_mode: SshMode = ...,
+            separate_stderr: bool = ...,
+            connect_timeout: Optional[int] = ...,
+            source_bashrc: bool = ...,
+            skip_lines: int = ...,
+            **kwargs) -> int:
+        ...
+
+    @typing.overload
+    def run(self,
+            cmd: Union[str, List[str]],
+            *,
+            port_forward: Optional[List[int]] = ...,
+            require_outputs: Literal[True],
+            log_path: str = ...,
+            process_stream: bool = ...,
+            stream_logs: bool = ...,
+            ssh_mode: SshMode = ...,
+            separate_stderr: bool = ...,
+            connect_timeout: Optional[int] = ...,
+            source_bashrc: bool = ...,
+            skip_lines: int = ...,
+            **kwargs) -> Tuple[int, str, str]:
+        ...
+
+    @typing.overload
+    def run(self,
+            cmd: Union[str, List[str]],
+            *,
+            port_forward: Optional[List[int]] = ...,
+            require_outputs: bool = ...,
+            log_path: str = ...,
+            process_stream: bool = ...,
+            stream_logs: bool = ...,
+            ssh_mode: SshMode = ...,
+            separate_stderr: bool = ...,
+            connect_timeout: Optional[int] = ...,
+            source_bashrc: bool = ...,
+            skip_lines: int = ...,
+            **kwargs) -> Union[Tuple[int, str, str], int]:
+        ...

sky/utils/common.py

@@ -5,30 +5,56 @@ import enum
 import os
 from typing import Generator
 
+from sky import models
 from sky.skylet import constants
 from sky.utils import common_utils
 
 SKY_SERVE_CONTROLLER_PREFIX: str = 'sky-serve-controller-'
 JOB_CONTROLLER_PREFIX: str = 'sky-jobs-controller-'
+
 # We use the user hash (machine-specific) for the controller name. It will be
 # the same across the whole lifecycle of the server, including:
-# 1. all requests, because this global variable is set once during server
-#    starts.
-# 2. SkyPilot API server restarts, as long as the `~/.sky` folder is persisted
-#    and the env var set during starting the server is the same.
+# 1. all requests, because all the server processes share the same user hash
+#    cache file.
+# 2. SkyPilot API server restarts, because the API server will restore the
+#    user hash from the global user state db on startup.
+# 3. Potential multiple server replicas, because multiple server replicas of
+#    a same deployment will share the same global user state db.
 # This behavior is the same for the local API server (where SERVER_ID is the
 # same as the normal user hash). This ensures backwards-compatibility with jobs
 # controllers from before #4660.
-SERVER_ID = common_utils.get_user_hash()
-SKY_SERVE_CONTROLLER_NAME: str = f'{SKY_SERVE_CONTROLLER_PREFIX}{SERVER_ID}'
-JOB_CONTROLLER_NAME: str = f'{JOB_CONTROLLER_PREFIX}{SERVER_ID}'
+SERVER_ID: str
+SKY_SERVE_CONTROLLER_NAME: str
+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 authoritative
+    source.
+    """
+    global SERVER_ID
+    global SKY_SERVE_CONTROLLER_NAME
+    global JOB_CONTROLLER_NAME
+    SERVER_ID = common_utils.get_user_hash()
+    SKY_SERVE_CONTROLLER_NAME = f'{SKY_SERVE_CONTROLLER_PREFIX}{SERVER_ID}'
+    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()
 
 
 @contextlib.contextmanager
-def with_server_user_hash() -> Generator[None, None, None]:
+def with_server_user() -> Generator[None, None, None]:
     """Temporarily set the user hash to common.SERVER_ID."""
     old_env_user_hash = os.getenv(constants.USER_ID_ENV_VAR)
+    # TODO(zhwu): once we have fully moved our code to use `get_current_user()`
+    # instead of `common_utils.get_user_hash()`, we can remove the env override.
     os.environ[constants.USER_ID_ENV_VAR] = SERVER_ID
+    common_utils.set_current_user(models.User.get_current_user())
     try:
         yield
     finally:
@@ -36,6 +62,7 @@ def with_server_user_hash() -> Generator[None, None, None]:
             os.environ[constants.USER_ID_ENV_VAR] = old_env_user_hash
         else:
             os.environ.pop(constants.USER_ID_ENV_VAR)
+        common_utils.set_current_user(models.User.get_current_user())
 
 
 class StatusRefreshMode(enum.Enum):