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

sky/utils/command_runner.pyi

@@ -6,7 +6,7 @@ determine the return type based on the value of require_outputs.
 """
 import enum
 import typing
-from typing import Any, Iterable, List, Optional, Tuple, Union
+from typing import Any, Callable, Iterable, List, Optional, Tuple, Union
 
 from typing_extensions import Literal
 
@@ -27,6 +27,7 @@ def ssh_options_list(
     ssh_control_name: Optional[str],
     *,
     ssh_proxy_command: Optional[str] = ...,
+    ssh_proxy_jump: Optional[str] = ...,
     docker_ssh_proxy_command: Optional[str] = ...,
     timeout: int = ...,
     port: int = ...,
@@ -63,6 +64,7 @@ class CommandRunner:
             connect_timeout: Optional[int] = ...,
             source_bashrc: bool = ...,
             skip_lines: int = ...,
+            run_in_background: bool = ...,
             **kwargs) -> int:
         ...
 
@@ -78,6 +80,7 @@ class CommandRunner:
             connect_timeout: Optional[int] = ...,
             source_bashrc: bool = ...,
             skip_lines: int = ...,
+            run_in_background: bool = ...,
             **kwargs) -> Tuple[int, str, str]:
         ...
 
@@ -93,6 +96,7 @@ class CommandRunner:
             connect_timeout: Optional[int] = ...,
             source_bashrc: bool = ...,
             skip_lines: int = ...,
+            run_in_background: bool = ...,
             **kwargs) -> Union[Tuple[int, str, str], int]:
         ...
 
@@ -130,22 +134,25 @@ class SSHCommandRunner(CommandRunner):
     ip: str
     port: int
     ssh_user: str
-    ssh_private_key: str
+    ssh_private_key: Optional[str]
     ssh_control_name: Optional[str]
     docker_user: str
     disable_control_master: Optional[bool]
     port_forward_execute_remote_command: Optional[bool]
+    enable_interactive_auth: bool
 
     def __init__(
         self,
         node: Tuple[str, int],
         ssh_user: str,
-        ssh_private_key: str,
+        ssh_private_key: Optional[str],
         ssh_control_name: Optional[str] = ...,
         ssh_proxy_command: Optional[str] = ...,
+        ssh_proxy_jump: Optional[str] = ...,
         docker_user: Optional[str] = ...,
         disable_control_master: Optional[bool] = ...,
         port_forward_execute_remote_command: Optional[bool] = ...,
+        enable_interactive_auth: bool = ...,
     ) -> None:
         ...
 
@@ -163,6 +170,7 @@ class SSHCommandRunner(CommandRunner):
             connect_timeout: Optional[int] = ...,
             source_bashrc: bool = ...,
             skip_lines: int = ...,
+            run_in_background: bool = ...,
             **kwargs) -> int:
         ...
 
@@ -180,6 +188,7 @@ class SSHCommandRunner(CommandRunner):
             connect_timeout: Optional[int] = ...,
             source_bashrc: bool = ...,
             skip_lines: int = ...,
+            run_in_background: bool = ...,
             **kwargs) -> Tuple[int, str, str]:
         ...
 
@@ -197,6 +206,7 @@ class SSHCommandRunner(CommandRunner):
             connect_timeout: Optional[int] = ...,
             source_bashrc: bool = ...,
             skip_lines: int = ...,
+            run_in_background: bool = ...,
             **kwargs) -> Union[Tuple[int, str, str], int]:
         ...
 
@@ -216,7 +226,8 @@ class SSHCommandRunner(CommandRunner):
               up: bool,
               log_path: str = ...,
               stream_logs: bool = ...,
-              max_retry: int = ...) -> None:
+              max_retry: int = ...,
+              get_remote_home_dir: Callable[[], str] = ...) -> None:
         ...
 
     def port_forward_command(
@@ -251,6 +262,7 @@ class KubernetesCommandRunner(CommandRunner):
             connect_timeout: Optional[int] = ...,
             source_bashrc: bool = ...,
             skip_lines: int = ...,
+            run_in_background: bool = ...,
             **kwargs) -> int:
         ...
 
@@ -268,6 +280,7 @@ class KubernetesCommandRunner(CommandRunner):
             connect_timeout: Optional[int] = ...,
             source_bashrc: bool = ...,
             skip_lines: int = ...,
+            run_in_background: bool = ...,
             **kwargs) -> Tuple[int, str, str]:
         ...
 
@@ -285,6 +298,7 @@ class KubernetesCommandRunner(CommandRunner):
             connect_timeout: Optional[int] = ...,
             source_bashrc: bool = ...,
             skip_lines: int = ...,
+            run_in_background: bool = ...,
             **kwargs) -> Union[Tuple[int, str, str], int]:
         ...
 
@@ -306,6 +320,28 @@ class KubernetesCommandRunner(CommandRunner):
         ...
 
 
+class SlurmCommandRunner(SSHCommandRunner):
+    """Runner for Slurm commands."""
+    sky_dir: str
+    skypilot_runtime_dir: str
+    job_id: str
+    slurm_node: str
+
+    def __init__(
+        self,
+        node: Tuple[str, int],
+        ssh_user: str,
+        ssh_private_key: Optional[str],
+        *,
+        sky_dir: str,
+        skypilot_runtime_dir: str,
+        job_id: str,
+        slurm_node: str,
+        **kwargs,
+    ) -> None:
+        ...
+
+
 class LocalProcessCommandRunner(CommandRunner):
 
     def __init__(self) -> None:
@@ -325,6 +361,7 @@ class LocalProcessCommandRunner(CommandRunner):
             connect_timeout: Optional[int] = ...,
             source_bashrc: bool = ...,
             skip_lines: int = ...,
+            run_in_background: bool = ...,
             **kwargs) -> int:
         ...
 
@@ -342,6 +379,7 @@ class LocalProcessCommandRunner(CommandRunner):
             connect_timeout: Optional[int] = ...,
             source_bashrc: bool = ...,
             skip_lines: int = ...,
+            run_in_background: bool = ...,
             **kwargs) -> Tuple[int, str, str]:
         ...
 
@@ -359,5 +397,6 @@ class LocalProcessCommandRunner(CommandRunner):
             connect_timeout: Optional[int] = ...,
             source_bashrc: bool = ...,
             skip_lines: int = ...,
+            run_in_background: bool = ...,
             **kwargs) -> Union[Tuple[int, str, str], int]:
         ...

sky/utils/common_utils.py

@@ -29,6 +29,7 @@ from sky.adaptors import common as adaptors_common
 from sky.skylet import constants
 from sky.usage import constants as usage_constants
 from sky.utils import annotations
+from sky.utils import context
 from sky.utils import ux_utils
 from sky.utils import validator
 
@@ -293,11 +294,11 @@ class Backoff:
         return self._backoff
 
 
-_current_command: Optional[str] = None
-_current_client_entrypoint: Optional[str] = None
-_using_remote_api_server: Optional[bool] = None
-_current_user: Optional['models.User'] = None
-_current_request_id: Optional[str] = None
+_CLIENT_COMMAND_KEY = 'client_command'
+_CLIENT_ENTRYPOINT_KEY = 'client_entrypoint'
+_USING_REMOTE_API_SERVER_KEY = 'using_remote_api_server'
+_USER_KEY = 'user'
+_REQUEST_ID_KEY = 'request_id'
 
 
 def set_request_context(client_entrypoint: Optional[str],
@@ -309,22 +310,21 @@ def set_request_context(client_entrypoint: Optional[str],
     This is useful when we are on the SkyPilot API server side and we have a
     client entrypoint and command from the client.
     """
-    global _current_command
-    global _current_client_entrypoint
-    global _using_remote_api_server
-    global _current_user
-    global _current_request_id
-    _current_command = client_command
-    _current_client_entrypoint = client_entrypoint
-    _using_remote_api_server = using_remote_api_server
-    _current_user = user
-    _current_request_id = request_id
+    # This function will be called in process executor and coroutine executor.
+    # context.set_context_var ensures the context is safe in both cases.
+    context.set_context_var(_CLIENT_ENTRYPOINT_KEY, client_entrypoint)
+    context.set_context_var(_CLIENT_COMMAND_KEY, client_command)
+    context.set_context_var(_USING_REMOTE_API_SERVER_KEY,
+                            using_remote_api_server)
+    context.set_context_var(_USER_KEY, user)
+    context.set_context_var(_REQUEST_ID_KEY, request_id)
 
 
 def get_current_request_id() -> str:
     """Returns the current request id."""
-    if _current_request_id is not None:
-        return _current_request_id
+    value = context.get_context_var('request_id')
+    if value is not None:
+        return value
     return 'dummy-request-id'
 
 
@@ -334,30 +334,43 @@ def get_current_command() -> str:
     Normally uses get_pretty_entry_point(), but will use the client command on
     the server side.
     """
-    if _current_command is not None:
-        return _current_command
-
+    value = context.get_context_var(_CLIENT_COMMAND_KEY)
+    if value is not None:
+        return value
     return get_pretty_entrypoint_cmd()
 
 
 def get_current_user() -> 'models.User':
-    """Returns the current user."""
-    if _current_user is not None:
-        return _current_user
+    """Returns the user in current server session."""
+    value = context.get_context_var(_USER_KEY)
+    if value is not None:
+        return value
     return models.User.get_current_user()
 
 
 def get_current_user_name() -> str:
-    """Returns the current user name."""
+    """Returns the user name in current server session."""
     name = get_current_user().name
     assert name is not None
     return name
 
 
+def get_local_user_name() -> str:
+    """Returns the user name in local environment.
+
+    This is for backward compatibility where anonymous access is implicitly
+    allowed when no authentication method at server-side is configured and
+    the username from client environment variable will be used to identify the
+    user.
+    """
+    name = os.getenv(constants.USER_ENV_VAR, getpass.getuser())
+    assert name is not None
+    return name
+
+
 def set_current_user(user: 'models.User'):
     """Sets the current user."""
-    global _current_user
-    _current_user = user
+    context.set_context_var('user', user)
 
 
 def get_current_client_entrypoint(server_entrypoint: str) -> str:
@@ -366,8 +379,9 @@ def get_current_client_entrypoint(server_entrypoint: str) -> str:
     Gets the client entrypoint from the context, if it is not set, returns the
     server entrypoint.
     """
-    if _current_client_entrypoint is not None:
-        return _current_client_entrypoint
+    value = context.get_context_var(_CLIENT_ENTRYPOINT_KEY)
+    if value is not None:
+        return value
     return server_entrypoint
 
 
@@ -376,8 +390,9 @@ def get_using_remote_api_server() -> bool:
     if os.getenv(constants.USING_REMOTE_API_SERVER_ENV_VAR) is not None:
         return os.getenv(constants.USING_REMOTE_API_SERVER_ENV_VAR,
                          '').lower() in ('true', '1')
-    if _using_remote_api_server is not None:
-        return _using_remote_api_server
+    value = context.get_context_var(_USING_REMOTE_API_SERVER_KEY)
+    if value is not None:
+        return value
     # This gets the right status for the local client.
     # TODO(zhwu): This is to prevent circular import. We should refactor this.
     # pylint: disable=import-outside-toplevel
@@ -724,7 +739,8 @@ def find_free_port(start_port: int) -> int:
             try:
                 s.bind(('', port))
                 return port
-            except OSError:
+            except OSError as e:
+                logger.debug(f'Error binding port {port}: {e}')
                 pass
     raise OSError('No free ports available.')
 

sky/utils/context.py

@@ -17,6 +17,8 @@ from typing_extensions import ParamSpec
 if TYPE_CHECKING:
     from sky.skypilot_config import ConfigContext
 
+_PROCESS_GLOBAL_VARS = {}
+
 
 class SkyPilotContext(object):
     """SkyPilot typed context vars for threads and coroutines.
@@ -65,6 +67,8 @@ class SkyPilotContext(object):
         self._log_file_handle = None
         self.env_overrides = {}
         self.config_context = None
+        self.request_context = None
+        self.vars = {}
 
     def cancel(self):
         """Cancel the context."""
@@ -113,6 +117,12 @@ class SkyPilotContext(object):
             self._log_file_handle.close()
             self._log_file_handle = None
 
+    def set_var(self, key: str, value: Any):
+        self.vars[key] = value
+
+    def get_var(self, key: str) -> Optional[Any]:
+        return self.vars.get(key)
+
     def __enter__(self):
         return self
 
@@ -150,6 +160,28 @@ def get() -> Optional[SkyPilotContext]:
     return _CONTEXT.get()
 
 
+def set_context_var(key: str, value: Any):
+    ctx = get()
+    if ctx is not None:
+        # Set the var in context
+        ctx.set_var(key, value)
+    else:
+        # Fallback to process-isolated assumption, where we thought
+        # modifying process-scope vars is safe.
+        _PROCESS_GLOBAL_VARS[key] = value
+
+
+def get_context_var(key: str) -> Any:
+    ctx = get()
+    if ctx is not None:
+        # Use `in` to check for key existence to distinguish
+        # "key not found" from "key's value is None".
+        if key in ctx.vars:
+            return ctx.get_var(key)
+    # Fallback to the variable set in process-scope
+    return _PROCESS_GLOBAL_VARS.get(key)
+
+
 class ContextualEnviron(MutableMapping[str, str]):
     """Environment variables wrapper with contextual overrides.
 

sky/utils/db/db_utils.py

@@ -75,6 +75,18 @@ def safe_cursor(db_path: str):
         conn.close()
 
 
+@contextlib.contextmanager
+def safe_cursor_on_connection(conn: 'sqlite3.Connection'):
+    """A auto-committing, auto-closing cursor on an existing connection."""
+    # Ensure commit() is called when the context is exited.
+    with conn:
+        cursor = conn.cursor()
+        try:
+            yield cursor
+        finally:
+            cursor.close()
+
+
 def add_column_to_table(
     cursor: 'sqlite3.Cursor',
     conn: 'sqlite3.Connection',
@@ -286,6 +298,11 @@ def drop_column_from_table_alembic(
             raise
 
 
+def fault_point():
+    """For test fault injection."""
+    pass
+
+
 class SQLiteConn(threading.local):
     """Thread-local connection to the sqlite3 database."""
 
@@ -345,8 +362,8 @@ class SQLiteConn(threading.local):
 
         def exec_and_commit(sql: str, parameters: Optional[Iterable[Any]]):
             # pylint: disable=protected-access
-            conn._conn.execute(sql, parameters)
-            conn._conn.commit()
+            with safe_cursor_on_connection(conn._conn) as cursor:
+                cursor.execute(sql, parameters)
 
         # pylint: disable=protected-access
         await conn._execute(exec_and_commit, sql, parameters)
@@ -357,7 +374,20 @@ class SQLiteConn(threading.local):
                                      parameters: Optional[Iterable[Any]] = None
                                     ) -> Iterable[sqlite3.Row]:
         conn = await self._get_async_conn()
-        return await conn.execute_fetchall(sql, parameters)
+        if parameters is None:
+            parameters = []
+
+        def exec_fetch_all(sql: str, parameters: Optional[Iterable[Any]]):
+            # pylint: disable=protected-access
+            with safe_cursor_on_connection(conn._conn) as cursor:
+                cursor.execute(sql, parameters)
+                # Note(dev): sqlite3.Connection cannot be patched, keep
+                # fault_point here to test the integrity of exec_fetch_all()
+                fault_point()
+                return cursor.fetchall()
+
+        # pylint: disable=protected-access
+        return await conn._execute(exec_fetch_all, sql, parameters)
 
     async def execute_get_returning_value_async(
             self,
@@ -372,9 +402,9 @@ class SQLiteConn(threading.local):
         def exec_and_get_returning_value(sql: str,
                                          parameters: Optional[Iterable[Any]]):
             # pylint: disable=protected-access
-            row = conn._conn.execute(sql, parameters).fetchone()
-            conn._conn.commit()
-            return row
+            with safe_cursor_on_connection(conn._conn) as cursor:
+                cursor.execute(sql, parameters)
+                return cursor.fetchone()
 
         # pylint: disable=protected-access
         return await conn._execute(exec_and_get_returning_value, sql,

sky/utils/db/migration_utils.py

@@ -3,6 +3,7 @@
 import contextlib
 import logging
 import os
+from typing import Optional
 
 from alembic import command as alembic_command
 from alembic.config import Config
@@ -22,7 +23,7 @@ GLOBAL_USER_STATE_VERSION = '011'
 GLOBAL_USER_STATE_LOCK_PATH = f'~/.sky/locks/.{GLOBAL_USER_STATE_DB_NAME}.lock'
 
 SPOT_JOBS_DB_NAME = 'spot_jobs_db'
-SPOT_JOBS_VERSION = '007'
+SPOT_JOBS_VERSION = '011'
 SPOT_JOBS_LOCK_PATH = f'~/.sky/locks/.{SPOT_JOBS_DB_NAME}.lock'
 
 SERVE_DB_NAME = 'serve_db'
@@ -52,12 +53,22 @@ def db_lock(db_name: str):
                            f'file if you believe it is stale.') from e
 
 
-def get_alembic_config(engine: sqlalchemy.engine.Engine, section: str):
-    """Get Alembic configuration for the given section"""
-    # From sky/utils/db/migration_utils.py -> sky/setup_files/alembic.ini
-    alembic_ini_path = os.path.join(
-        os.path.dirname(os.path.dirname(os.path.dirname(__file__))),
-        'setup_files', 'alembic.ini')
+def get_alembic_config(engine: sqlalchemy.engine.Engine,
+                       section: str,
+                       alembic_ini_path: Optional[str] = None):
+    """Get Alembic configuration for the given section.
+
+    Args:
+        engine: SQLAlchemy engine for the database.
+        section: Alembic section name (e.g., 'state_db' or 'spot_jobs_db').
+        alembic_ini_path: Optional path to a custom alembic.ini file.
+            If not provided, uses the default SkyPilot alembic.ini.
+    """
+    if alembic_ini_path is None:
+        # Default to SkyPilot's alembic.ini
+        alembic_ini_path = os.path.join(
+            os.path.dirname(os.path.dirname(os.path.dirname(__file__))),
+            'setup_files', 'alembic.ini')
     alembic_cfg = Config(alembic_ini_path, ini_section=section)
 
     # Override the database URL to match SkyPilot's current connection
@@ -73,19 +84,23 @@ def get_alembic_config(engine: sqlalchemy.engine.Engine, section: str):
     return alembic_cfg
 
 
-def needs_upgrade(engine: sqlalchemy.engine.Engine, section: str,
-                  target_revision: str):
+def needs_upgrade(engine: sqlalchemy.engine.Engine,
+                  section: str,
+                  target_revision: str,
+                  alembic_ini_path: Optional[str] = None):
     """Check if the database needs to be upgraded.
 
     Args:
-        engine: SQLAlchemy engine for the database
-        section: Alembic section to upgrade (e.g., 'state_db' or 'spot_jobs_db')
-        target_revision: Target revision to upgrade to (e.g., '001')
+        engine: SQLAlchemy engine for the database.
+        section: Alembic section to upgrade (e.g., 'state_db' or
+        'spot_jobs_db').
+        target_revision: Target revision to upgrade to (e.g., '001').
+        alembic_ini_path: Optional path to a custom alembic.ini file.
     """
     current_rev = None
 
     # get alembic config for the given section
-    alembic_config = get_alembic_config(engine, section)
+    alembic_config = get_alembic_config(engine, section, alembic_ini_path)
     version_table = alembic_config.get_section_option(
         alembic_config.config_ini_section, 'version_table', 'alembic_version')
 
@@ -112,26 +127,31 @@ def needs_upgrade(engine: sqlalchemy.engine.Engine, section: str,
     return current_rev_num < target_rev_num
 
 
-def safe_alembic_upgrade(engine: sqlalchemy.engine.Engine, section: str,
-                         target_revision: str):
+def safe_alembic_upgrade(engine: sqlalchemy.engine.Engine,
+                         section: str,
+                         target_revision: str,
+                         alembic_ini_path: Optional[str] = None):
     """Upgrade the database if needed. Uses a file lock to ensure
     that only one process tries to upgrade the database at a time.
 
     Args:
-        engine: SQLAlchemy engine for the database
-        section: Alembic section to upgrade (e.g., 'state_db' or 'spot_jobs_db')
-        target_revision: Target revision to upgrade to (e.g., '001')
+        engine: SQLAlchemy engine for the database.
+        section: Alembic section to upgrade (e.g., 'state_db' or
+        'spot_jobs_db').
+        target_revision: Target revision to upgrade to (e.g., '001').
+        alembic_ini_path: Optional path to a custom alembic.ini file.
     """
     # set alembic logger to warning level
     alembic_logger = logging.getLogger('alembic')
     alembic_logger.setLevel(logging.WARNING)
 
-    alembic_config = get_alembic_config(engine, section)
+    alembic_config = get_alembic_config(engine, section, alembic_ini_path)
 
     # only acquire lock if db needs upgrade
-    if needs_upgrade(engine, section, target_revision):
+    if needs_upgrade(engine, section, target_revision, alembic_ini_path):
         with db_lock(section):
             # check again if db needs upgrade in case another
             # process upgraded it while we were waiting for the lock
-            if needs_upgrade(engine, section, target_revision):
+            if needs_upgrade(engine, section, target_revision,
+                             alembic_ini_path):
                 alembic_command.upgrade(alembic_config, target_revision)

sky/utils/infra_utils.py

@@ -173,7 +173,11 @@ class InfraInfo:
             return '-'
 
         region_or_zone = None
-        if self.zone is not None and self.zone != '*':
+        # For Slurm, zones = partitions. We want to show the cluster
+        # name (region) instead of the partition name (zone), as different
+        # Slurm clusters can easily have same partition name.
+        is_slurm = self.cloud.lower() == 'slurm'
+        if not is_slurm and self.zone is not None and self.zone != '*':
             region_or_zone = self.zone
         elif self.region is not None and self.region != '*':
             # If using region, we remove the ssh- prefix if it exists for SSH