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

sky/skylet/subprocess_daemon.py

@@ -4,11 +4,16 @@ processes of proc_pid.
 """
 import argparse
 import os
+import signal
 import sys
 import time
+from typing import List, Optional
 
 import psutil
 
+# Environment variable to enable kill_pg in subprocess daemon.
+USE_KILL_PG_ENV_VAR = 'SKYPILOT_SUBPROCESS_DAEMON_KILL_PG'
+
 
 def daemonize():
     """Detaches the process from its parent process with double-forking.
@@ -38,8 +43,75 @@ def daemonize():
     # This process is now fully detached from the original parent and terminal
 
 
-if __name__ == '__main__':
+def get_pgid_if_leader(pid) -> Optional[int]:
+    """Get the process group ID of the target process if it is the leader."""
+    try:
+        pgid = os.getpgid(pid)
+        # Only use process group if the target process is the leader. This is
+        # to avoid killing the entire process group while the target process is
+        # just a subprocess in the group.
+        if pgid == pid:
+            print(f'Process group {pgid} is the leader.')
+            return pgid
+        return None
+    except Exception:  # pylint: disable=broad-except
+        # Process group is only available in UNIX.
+        return None
+
+
+def kill_process_group(pgid: int) -> bool:
+    """Kill the target process group."""
+    try:
+        print(f'Terminating process group {pgid}...')
+        os.killpg(pgid, signal.SIGTERM)
+    except Exception:  # pylint: disable=broad-except
+        return False
+
+    # Wait 30s for the process group to exit gracefully.
+    time.sleep(30)
+
+    try:
+        print(f'Force killing process group {pgid}...')
+        os.killpg(pgid, signal.SIGKILL)
+    except Exception:  # pylint: disable=broad-except
+        pass
+
+    return True
+
+
+def kill_process_tree(process: psutil.Process,
+                      children: List[psutil.Process]) -> bool:
+    """Kill the process tree of the target process."""
+    if process is not None:
+        # Kill the target process first to avoid having more children, or fail
+        # the process due to the children being defunct.
+        children = [process] + children
+
+    if not children:
+        sys.exit()
+
+    for child in children:
+        try:
+            child.terminate()
+        except psutil.NoSuchProcess:
+            continue
+
+    # Wait 30s for the processes to exit gracefully.
+    time.sleep(30)
+
+    # SIGKILL if they're still running.
+    for child in children:
+        try:
+            child.kill()
+        except psutil.NoSuchProcess:
+            continue
+
+    return True
+
+
+def main():
     daemonize()
+
     parser = argparse.ArgumentParser()
     parser.add_argument('--parent-pid', type=int, required=True)
     parser.add_argument('--proc-pid', type=int, required=True)
@@ -72,37 +144,41 @@ if __name__ == '__main__':
             except (psutil.NoSuchProcess, ValueError):
                 pass
 
+    pgid: Optional[int] = None
+    if os.environ.get(USE_KILL_PG_ENV_VAR) == '1':
+        # Use kill_pg on UNIX system if allowed to reduce the resource usage.
+        # Note that both implementations might leave subprocessed uncancelled:
+        # - kill_process_tree(default): a subprocess is able to detach itself
+        #   from the process tree use the same technique as daemonize(). Also,
+        #   since we refresh the process tree per second, if the subprocess is
+        #   launched between the [last_poll, parent_die] interval, the
+        #   subprocess will not be captured will not be killed.
+        # - kill_process_group: kill_pg will kill all the processed in the group
+        #   but if a subprocess calls setpgid(0, 0) to detach itself from the
+        #   process group (usually to daemonize itself), the subprocess will
+        #   not be killed.
+        if process is not None:
+            pgid = get_pgid_if_leader(process.pid)
+
     if process is not None and parent_process is not None:
         # Wait for either parent or target process to exit
         while process.is_running() and parent_process.is_running():
-            try:
-                tmp_children = process.children(recursive=True)
-                if tmp_children:
-                    children = tmp_children
-            except psutil.NoSuchProcess:
-                pass
+            if pgid is None:
+                # Refresh process tree for cleanup if process group is not
+                # available.
+                try:
+                    tmp_children = process.children(recursive=True)
+                    if tmp_children:
+                        children = tmp_children
+                except psutil.NoSuchProcess:
+                    pass
             time.sleep(1)
 
-    if process is not None:
-        # Kill the target process first to avoid having more children, or fail
-        # the process due to the children being defunct.
-        children = [process] + children
-
-    if not children:
-        sys.exit()
-
-    for child in children:
-        try:
-            child.terminate()
-        except psutil.NoSuchProcess:
-            continue
+    if pgid is not None:
+        kill_process_group(pgid)
+    else:
+        kill_process_tree(process, children)
 
-    # Wait 30s for the processes to exit gracefully.
-    time.sleep(30)
 
-    # SIGKILL if they're still running.
-    for child in children:
-        try:
-            child.kill()
-        except psutil.NoSuchProcess:
-            continue
+if __name__ == '__main__':
+    main()

sky/skypilot_config.py

@@ -64,7 +64,6 @@ from sqlalchemy import orm
 from sqlalchemy.dialects import postgresql
 from sqlalchemy.dialects import sqlite
 from sqlalchemy.ext import declarative
-from sqlalchemy.pool import NullPool
 
 from sky import exceptions
 from sky import sky_logging
@@ -77,6 +76,7 @@ from sky.utils import schemas
 from sky.utils import ux_utils
 from sky.utils import yaml_utils
 from sky.utils.db import db_utils
+from sky.utils.db import migration_utils
 from sky.utils.kubernetes import config_map_utils
 
 if typing.TYPE_CHECKING:
@@ -121,7 +121,8 @@ _PROJECT_CONFIG_PATH = '.sky.yaml'
 
 API_SERVER_CONFIG_KEY = 'api_server_config'
 
-_DB_USE_LOCK = threading.Lock()
+_SQLALCHEMY_ENGINE: Optional[sqlalchemy.engine.Engine] = None
+_SQLALCHEMY_ENGINE_LOCK = threading.Lock()
 
 Base = declarative.declarative_base()
 
@@ -415,10 +416,17 @@ def local_active_workspace_ctx(workspace: str) -> Iterator[None]:
 def get_active_workspace(force_user_workspace: bool = False) -> str:
     context_workspace = getattr(_active_workspace_context, 'workspace', None)
     if not force_user_workspace and context_workspace is not None:
-        logger.debug(f'Get context workspace: {context_workspace}')
+        logger.debug(f'Got context workspace: {context_workspace}')
         return context_workspace
-    return get_nested(keys=('active_workspace',),
-                      default_value=constants.SKYPILOT_DEFAULT_WORKSPACE)
+    active_workspace = get_nested(keys=('active_workspace',),
+                                  default_value=None)
+    if active_workspace is None:
+        logger.debug(f'No active workspace found, using default workspace: '
+                     f'{constants.SKYPILOT_DEFAULT_WORKSPACE}')
+        active_workspace = constants.SKYPILOT_DEFAULT_WORKSPACE
+    else:
+        logger.debug(f'Got active workspace: {active_workspace}')
+    return active_workspace
 
 
 def set_nested(keys: Tuple[str, ...], value: Any) -> Dict[str, Any]:
@@ -474,7 +482,7 @@ def safe_reload_config() -> None:
         reload_config()
 
 
-def reload_config() -> None:
+def reload_config(init_db: bool = False) -> None:
     internal_config_path = os.environ.get(ENV_VAR_SKYPILOT_CONFIG)
     if internal_config_path is not None:
         # {ENV_VAR_SKYPILOT_CONFIG} is used internally.
@@ -486,7 +494,7 @@ def reload_config() -> None:
         return
 
     if os.environ.get(constants.ENV_VAR_IS_SKYPILOT_SERVER) is not None:
-        _reload_config_as_server()
+        _reload_config_as_server(init_db=init_db)
     else:
         _reload_config_as_client()
 
@@ -557,7 +565,43 @@ def _reload_config_from_internal_file(internal_config_path: str) -> None:
     _set_loaded_config_path(config_path)
 
 
-def _reload_config_as_server() -> None:
+def _create_table(engine: sqlalchemy.engine.Engine):
+    """Initialize the config database with migrations."""
+    migration_utils.safe_alembic_upgrade(
+        engine, migration_utils.SKYPILOT_CONFIG_DB_NAME,
+        migration_utils.SKYPILOT_CONFIG_VERSION)
+
+
+def _initialize_and_get_db() -> sqlalchemy.engine.Engine:
+    """Initialize and return the config database engine.
+
+    This function should only be called by the API Server during initialization.
+    Client-side code should never call this function.
+    """
+    assert os.environ.get(constants.ENV_VAR_IS_SKYPILOT_SERVER) is not None, (
+        'initialize_and_get_db() can only be called by the API Server')
+
+    global _SQLALCHEMY_ENGINE
+
+    if _SQLALCHEMY_ENGINE is not None:
+        return _SQLALCHEMY_ENGINE
+
+    with _SQLALCHEMY_ENGINE_LOCK:
+        if _SQLALCHEMY_ENGINE is not None:
+            return _SQLALCHEMY_ENGINE
+
+        # We only store config in the DB when using Postgres,
+        # so no need to pass in db_name here.
+        engine = db_utils.get_engine(None)
+
+        # Run migrations if needed
+        _create_table(engine)
+
+        _SQLALCHEMY_ENGINE = engine
+        return _SQLALCHEMY_ENGINE
+
+
+def _reload_config_as_server(init_db: bool = False) -> None:
     # Reset the global variables, to avoid using stale values.
     _set_loaded_config(config_utils.Config())
     _set_loaded_config_path(None)
@@ -573,37 +617,24 @@ def _reload_config_as_server() -> None:
             raise ValueError(
                 'If db config is specified, no other config is allowed')
         logger.debug('retrieving config from database')
-        with _DB_USE_LOCK:
-            dispose_engine = False
-            if db_utils.get_max_connections() == 0:
-                dispose_engine = True
-                sqlalchemy_engine = sqlalchemy.create_engine(db_url,
-                                                             poolclass=NullPool)
-            else:
-                sqlalchemy_engine = db_utils.get_engine('config')
-            db_utils.add_all_tables_to_db_sqlalchemy(Base.metadata,
-                                                     sqlalchemy_engine)
-
-            def _get_config_yaml_from_db(
-                    key: str) -> Optional[config_utils.Config]:
-                assert sqlalchemy_engine is not None
-                with orm.Session(sqlalchemy_engine) as session:
-                    row = session.query(config_yaml_table).filter_by(
-                        key=key).first()
-                if row:
-                    db_config = config_utils.Config(
-                        yaml_utils.safe_load(row.value))
-                    db_config.pop_nested(('db',), None)
-                    return db_config
-                return None
-
-            db_config = _get_config_yaml_from_db(API_SERVER_CONFIG_KEY)
-            if db_config:
-                server_config = overlay_skypilot_config(server_config,
-                                                        db_config)
-            # Close the engine to avoid connection leaks
-            if dispose_engine:
-                sqlalchemy_engine.dispose()
+
+        if init_db:
+            _initialize_and_get_db()
+
+        def _get_config_yaml_from_db(key: str) -> Optional[config_utils.Config]:
+            assert _SQLALCHEMY_ENGINE is not None
+            with orm.Session(_SQLALCHEMY_ENGINE) as session:
+                row = session.query(config_yaml_table).filter_by(
+                    key=key).first()
+            if row:
+                db_config = config_utils.Config(yaml_utils.safe_load(row.value))
+                db_config.pop_nested(('db',), None)
+                return db_config
+            return None
+
+        db_config = _get_config_yaml_from_db(API_SERVER_CONFIG_KEY)
+        if db_config:
+            server_config = overlay_skypilot_config(server_config, db_config)
     if sky_logging.logging_enabled(logger, sky_logging.DEBUG):
         logger.debug(f'server config: \n'
                      f'{yaml_utils.dump_yaml_str(dict(server_config))}')
@@ -659,7 +690,7 @@ def loaded_config_path_serialized() -> Optional[str]:
 
 
 # Load on import, synchronization is guaranteed by python interpreter.
-reload_config()
+reload_config(init_db=True)
 
 
 def loaded() -> bool:
@@ -818,7 +849,8 @@ def _compose_cli_config(cli_config: Optional[List[str]]) -> config_utils.Config:
     except ValueError as e:
         raise ValueError(f'Invalid config override: {cli_config}. '
                          f'Check if config file exists or if the dotlist '
-                         f'is formatted as: key1=value1,key2=value2') from e
+                         f'is formatted as: key1=value1,key2=value2.\n'
+                         f'Details: {e}') from e
     logger.debug('CLI overrides config syntax check passed.')
 
     return parsed_config
@@ -872,44 +904,32 @@ def update_api_server_config_no_lock(config: config_utils.Config) -> None:
         if new_db_url and new_db_url != existing_db_url:
             raise ValueError('Cannot change db url while server is running')
         if existing_db_url:
-            with _DB_USE_LOCK:
-                dispose_engine = False
-                if db_utils.get_max_connections() == 0:
-                    dispose_engine = True
-                    sqlalchemy_engine = sqlalchemy.create_engine(
-                        existing_db_url, poolclass=NullPool)
-                else:
-                    sqlalchemy_engine = db_utils.get_engine('config')
-                db_utils.add_all_tables_to_db_sqlalchemy(
-                    Base.metadata, sqlalchemy_engine)
-
-                def _set_config_yaml_to_db(key: str,
-                                           config: config_utils.Config):
-                    assert sqlalchemy_engine is not None
-                    config_str = yaml_utils.dump_yaml_str(dict(config))
-                    with orm.Session(sqlalchemy_engine) as session:
-                        if (sqlalchemy_engine.dialect.name ==
-                                db_utils.SQLAlchemyDialect.SQLITE.value):
-                            insert_func = sqlite.insert
-                        elif (sqlalchemy_engine.dialect.name ==
-                              db_utils.SQLAlchemyDialect.POSTGRESQL.value):
-                            insert_func = postgresql.insert
-                        else:
-                            raise ValueError('Unsupported database dialect')
-                        insert_stmnt = insert_func(config_yaml_table).values(
-                            key=key, value=config_str)
-                        do_update_stmt = insert_stmnt.on_conflict_do_update(
-                            index_elements=[config_yaml_table.c.key],
-                            set_={config_yaml_table.c.value: config_str})
-                        session.execute(do_update_stmt)
-                        session.commit()
-
-                logger.debug('saving api_server config to db')
-                _set_config_yaml_to_db(API_SERVER_CONFIG_KEY, config)
-                db_updated = True
-                # Close the engine to avoid connection leaks
-                if dispose_engine:
-                    sqlalchemy_engine.dispose()
+
+            def _set_config_yaml_to_db(key: str, config: config_utils.Config):
+                # reload_config(init_db=True) is called when this module is
+                # imported, so the database engine must already be initialized.
+                assert _SQLALCHEMY_ENGINE is not None
+                config_str = yaml_utils.dump_yaml_str(dict(config))
+                with orm.Session(_SQLALCHEMY_ENGINE) as session:
+                    if (_SQLALCHEMY_ENGINE.dialect.name ==
+                            db_utils.SQLAlchemyDialect.SQLITE.value):
+                        insert_func = sqlite.insert
+                    elif (_SQLALCHEMY_ENGINE.dialect.name ==
+                          db_utils.SQLAlchemyDialect.POSTGRESQL.value):
+                        insert_func = postgresql.insert
+                    else:
+                        raise ValueError('Unsupported database dialect')
+                    insert_stmnt = insert_func(config_yaml_table).values(
+                        key=key, value=config_str)
+                    do_update_stmt = insert_stmnt.on_conflict_do_update(
+                        index_elements=[config_yaml_table.c.key],
+                        set_={config_yaml_table.c.value: config_str})
+                    session.execute(do_update_stmt)
+                    session.commit()
+
+            logger.debug('saving api_server config to db')
+            _set_config_yaml_to_db(API_SERVER_CONFIG_KEY, config)
+            db_updated = True
 
     if not db_updated:
         # save to the local file (PVC in Kubernetes, local file otherwise)

sky/ssh_node_pools/constants.py

@@ -0,0 +1,12 @@
+"""Constants for SSH Node Pools"""
+# pylint: disable=line-too-long
+import os
+
+DEFAULT_KUBECONFIG_PATH = os.path.expanduser('~/.kube/config')
+SSH_CONFIG_PATH = os.path.expanduser('~/.ssh/config')
+NODE_POOLS_INFO_DIR = os.path.expanduser('~/.sky/ssh_node_pools_info')
+NODE_POOLS_KEY_DIR = os.path.expanduser('~/.sky/ssh_keys')
+DEFAULT_SSH_NODE_POOLS_PATH = os.path.expanduser('~/.sky/ssh_node_pools.yaml')
+
+# TODO (kyuds): make this configurable?
+K3S_TOKEN = 'mytoken'  # Any string can be used as the token

sky/ssh_node_pools/core.py

@@ -1,10 +1,15 @@
 """SSH Node Pool management core functionality."""
 import os
 from pathlib import Path
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional, Tuple
 
 import yaml
 
+from sky import clouds
+from sky.ssh_node_pools import constants
+from sky.ssh_node_pools import deploy
+from sky.usage import usage_lib
+from sky.utils import common_utils
 from sky.utils import yaml_utils
 
 
@@ -12,8 +17,8 @@ class SSHNodePoolManager:
     """Manager for SSH Node Pool configurations."""
 
     def __init__(self):
-        self.config_path = Path.home() / '.sky' / 'ssh_node_pools.yaml'
-        self.keys_dir = Path.home() / '.sky' / 'ssh_keys'
+        self.config_path = Path(constants.DEFAULT_SSH_NODE_POOLS_PATH)
+        self.keys_dir = Path(constants.NODE_POOLS_KEY_DIR)
         self.keys_dir.mkdir(parents=True, exist_ok=True)
 
     def get_all_pools(self) -> Dict[str, Any]:
@@ -133,3 +138,35 @@ def list_ssh_keys() -> List[str]:
     """List available SSH keys."""
     manager = SSHNodePoolManager()
     return manager.list_ssh_keys()
+
+
+@usage_lib.entrypoint
+def ssh_up(infra: Optional[str] = None, cleanup: bool = False) -> None:
+    """Deploys or tears down a Kubernetes cluster on SSH targets.
+
+    Args:
+        infra: Name of the cluster configuration in ssh_node_pools.yaml.
+            If None, the first cluster in the file is used.
+        cleanup: If True, clean up the cluster instead of deploying.
+    """
+    deploy.run(cleanup=cleanup, infra=infra)
+
+
+@usage_lib.entrypoint
+def ssh_status(context_name: str) -> Tuple[bool, str]:
+    """Check the status of an SSH Node Pool context.
+
+    Args:
+        context_name: The SSH context name (e.g., 'ssh-my-cluster')
+
+    Returns:
+        Tuple[bool, str]: (is_ready, reason)
+            - is_ready: True if the SSH Node Pool is ready, False otherwise
+            - reason: Explanation of the status
+    """
+    try:
+        is_ready, reason = clouds.SSH.check_single_context(context_name)
+        return is_ready, reason
+    except Exception as e:  # pylint: disable=broad-except
+        return False, ('Failed to check SSH context: '
+                       f'{common_utils.format_exception(e)}')

sky/ssh_node_pools/deploy/__init__.py

@@ -0,0 +1,4 @@
+"""Module for Deploying SSH Node Pools"""
+from sky.ssh_node_pools.deploy.deploy import run
+
+__all__ = ['run']