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

sky/skylet/constants.py

@@ -100,7 +100,7 @@ TASK_ID_LIST_ENV_VAR = f'{SKYPILOT_ENV_VAR_PREFIX}TASK_IDS'
 # cluster yaml is updated.
 #
 # TODO(zongheng,zhanghao): make the upgrading of skylet automatic?
-SKYLET_VERSION = '22'
+SKYLET_VERSION = '25'
 # The version of the lib files that skylet/jobs use. Whenever there is an API
 # change for the job_lib or log_lib, we need to bump this version, so that the
 # user can be notified to update their SkyPilot version on the remote cluster.
@@ -226,7 +226,9 @@ RAY_INSTALLATION_COMMANDS = (
     f'{SKY_UV_PIP_CMD} list | grep "ray " | '
     f'grep {SKY_REMOTE_RAY_VERSION} 2>&1 > /dev/null '
     f'|| {RAY_STATUS} || '
-    f'{SKY_UV_PIP_CMD} install -U ray[default]=={SKY_REMOTE_RAY_VERSION}; '  # pylint: disable=line-too-long
+    # The pydantic-core==2.41.3 for arm seems corrupted
+    # so we need to avoid that specific version.
+    f'{SKY_UV_PIP_CMD} install -U "ray[default]=={SKY_REMOTE_RAY_VERSION}" "pydantic-core==2.41.1"; '  # pylint: disable=line-too-long
     # In some envs, e.g. pip does not have permission to write under /opt/conda
     # ray package will be installed under ~/.local/bin. If the user's PATH does
     # not include ~/.local/bin (the pip install will have the output: `WARNING:
@@ -402,10 +404,27 @@ OVERRIDEABLE_CONFIG_KEYS_IN_TASK: List[Tuple[str, ...]] = [
 ]
 # When overriding the SkyPilot configs on the API server with the client one,
 # we skip the following keys because they are meant to be client-side configs.
-SKIPPED_CLIENT_OVERRIDE_KEYS: List[Tuple[str, ...]] = [('api_server',),
-                                                       ('allowed_clouds',),
-                                                       ('workspaces',), ('db',),
-                                                       ('daemons',)]
+# Also, we skip the consolidation mode config as those should be only set on
+# the API server side.
+SKIPPED_CLIENT_OVERRIDE_KEYS: List[Tuple[str, ...]] = [
+    ('api_server',),
+    ('allowed_clouds',),
+    ('workspaces',),
+    ('db',),
+    ('daemons',),
+    # TODO(kevin,tian): Override the whole controller config once our test
+    # infrastructure supports setting dynamic server side configs.
+    # Tests that are affected:
+    # - test_managed_jobs_ha_kill_starting
+    # - test_managed_jobs_ha_kill_running
+    # - all tests that use LOW_CONTROLLER_RESOURCE_ENV or
+    #   LOW_CONTROLLER_RESOURCE_OVERRIDE_CONFIG (won't cause test failure,
+    #   but the configs won't be applied)
+    ('jobs', 'controller', 'consolidation_mode'),
+    ('serve', 'controller', 'consolidation_mode'),
+    ('jobs', 'controller', 'controller_logs_gc_retention_hours'),
+    ('jobs', 'controller', 'task_logs_gc_retention_hours'),
+]
 
 # Constants for Azure blob storage
 WAIT_FOR_STORAGE_ACCOUNT_CREATION = 60
@@ -455,6 +474,7 @@ ENV_VAR_DB_CONNECTION_URI = (f'{SKYPILOT_ENV_VAR_PREFIX}DB_CONNECTION_URI')
 # authentication is enabled in the API server.
 ENV_VAR_ENABLE_BASIC_AUTH = 'ENABLE_BASIC_AUTH'
 SKYPILOT_INITIAL_BASIC_AUTH = 'SKYPILOT_INITIAL_BASIC_AUTH'
+SKYPILOT_INGRESS_BASIC_AUTH_ENABLED = 'SKYPILOT_INGRESS_BASIC_AUTH_ENABLED'
 ENV_VAR_ENABLE_SERVICE_ACCOUNTS = 'ENABLE_SERVICE_ACCOUNTS'
 
 # Enable debug logging for requests.
@@ -471,7 +491,7 @@ CATALOG_DIR = '~/.sky/catalogs'
 ALL_CLOUDS = ('aws', 'azure', 'gcp', 'ibm', 'lambda', 'scp', 'oci',
               'kubernetes', 'runpod', 'vast', 'vsphere', 'cudo', 'fluidstack',
               'paperspace', 'primeintellect', 'do', 'nebius', 'ssh',
-              'hyperbolic', 'seeweb')
+              'hyperbolic', 'seeweb', 'shadeform')
 # END constants used for service catalog.
 
 # The user ID of the SkyPilot system.
@@ -531,3 +551,6 @@ ENV_VAR_LOOP_LAG_THRESHOLD_MS = (SKYPILOT_ENV_VAR_PREFIX +
 
 ARM64_ARCH = 'arm64'
 X86_64_ARCH = 'x86_64'
+
+SSH_DISABLE_LATENCY_MEASUREMENT_ENV_VAR = (
+    f'{SKYPILOT_ENV_VAR_PREFIX}SSH_DISABLE_LATENCY_MEASUREMENT')

sky/skylet/events.py

@@ -326,8 +326,15 @@ class AutostopEvent(SkyletEvent):
         cluster_name_on_cloud = cluster_config['cluster_name']
         is_cluster_multinode = cluster_config['max_workers'] > 0
 
+        # Clear AWS credentials from environment to force boto3 to use IAM
+        # role attached to the instance (lowest priority in credential chain).
+        # This allows the cluster to stop/terminate itself using its IAM role.
         os.environ.pop('AWS_ACCESS_KEY_ID', None)
         os.environ.pop('AWS_SECRET_ACCESS_KEY', None)
+        os.environ.pop('AWS_SESSION_TOKEN', None)
+        # Point boto3 to /dev/null to skip reading credentials from files.
+        os.environ['AWS_SHARED_CREDENTIALS_FILE'] = '/dev/null'
+        os.environ['AWS_CONFIG_FILE'] = '/dev/null'
 
         # Stop the ray autoscaler to avoid scaling up, during
         # stopping/terminating of the cluster.

sky/skylet/log_lib.py

@@ -220,7 +220,14 @@ def run_with_log(
                           stdin=stdin,
                           **kwargs) as proc:
         try:
-            subprocess_utils.kill_process_daemon(proc.pid)
+            if ctx is not None:
+                # When runs in coroutine, use kill_pg if available to avoid
+                # the overhead of refreshing the process tree in the daemon.
+                subprocess_utils.kill_process_daemon(proc.pid, use_kill_pg=True)
+            else:
+                # For backward compatibility, do not specify use_kill_pg by
+                # default.
+                subprocess_utils.kill_process_daemon(proc.pid)
             stdout = ''
             stderr = ''
             stdout_stream_handler = None
@@ -271,7 +278,6 @@ def run_with_log(
                 stdout, stderr = context_utils.pipe_and_wait_process(
                     ctx,
                     proc,
-                    cancel_callback=subprocess_utils.kill_children_processes,
                     stdout_stream_handler=stdout_stream_handler,
                     stderr_stream_handler=stderr_stream_handler)
             elif process_stream:

sky/skylet/log_lib.pyi

@@ -42,7 +42,7 @@ class _ProcessingArgs:
         ...
 
 
-def _get_context() -> Optional[context.Context]:
+def _get_context() -> Optional[context.SkyPilotContext]:
     ...
 
 

sky/skylet/services.py

@@ -216,10 +216,12 @@ class JobsServiceImpl(jobsv1_pb2_grpc.JobsServiceServicer):
                 if pool is not None:
                     pool_hash = serve_state.get_service_hash(pool)
                 # Add the managed job to job queue database.
+                user_id = managed_job.user_id if managed_job.HasField(
+                    'user_id') else None
                 managed_job_state.set_job_info(job_id, managed_job.name,
                                                managed_job.workspace,
                                                managed_job.entrypoint, pool,
-                                               pool_hash)
+                                               pool_hash, user_id)
                 # Set the managed job to PENDING state to make sure that
                 # this managed job appears in the `sky jobs queue`, even
                 # if it needs to wait to be submitted.
@@ -405,18 +407,22 @@ class ManagedJobsServiceImpl(managed_jobsv1_pb2_grpc.ManagedJobsServiceServicer
         context: grpc.ServicerContext
     ) -> managed_jobsv1_pb2.GetJobTableResponse:
         try:
-            accessible_workspaces = list(request.accessible_workspaces)
-            job_ids = list(request.job_ids.ids) if request.job_ids else None
+            accessible_workspaces = (
+                list(request.accessible_workspaces.workspaces)
+                if request.HasField('accessible_workspaces') else None)
+            job_ids = (list(request.job_ids.ids)
+                       if request.HasField('job_ids') else None)
             user_hashes: Optional[List[Optional[str]]] = None
-            if request.user_hashes:
+            if request.HasField('user_hashes'):
                 user_hashes = list(request.user_hashes.hashes)
                 # For backwards compatibility, we show jobs that do not have a
                 # user_hash. TODO: Remove before 0.12.0.
                 if request.show_jobs_without_user_hash:
                     user_hashes.append(None)
-            statuses = list(
-                request.statuses.statuses) if request.statuses else None
-
+            statuses = (list(request.statuses.statuses)
+                        if request.HasField('statuses') else None)
+            fields = (list(request.fields.fields)
+                      if request.HasField('fields') else None)
             job_queue = managed_job_utils.get_managed_job_queue(
                 skip_finished=request.skip_finished,
                 accessible_workspaces=accessible_workspaces,
@@ -430,7 +436,9 @@ class ManagedJobsServiceImpl(managed_jobsv1_pb2_grpc.ManagedJobsServiceServicer
                 page=request.page if request.HasField('page') else None,
                 limit=request.limit if request.HasField('limit') else None,
                 user_hashes=user_hashes,
-                statuses=statuses)
+                statuses=statuses,
+                fields=fields,
+            )
             jobs = job_queue['jobs']
             total = job_queue['total']
             total_no_filter = job_queue['total_no_filter']
@@ -438,7 +446,16 @@ class ManagedJobsServiceImpl(managed_jobsv1_pb2_grpc.ManagedJobsServiceServicer
 
             jobs_info = []
             for job in jobs:
+                converted_metadata = None
+                metadata = job.get('metadata')
+                if metadata:
+                    converted_metadata = {
+                        k: v for k, v in metadata.items() if v is not None
+                    }
                 job_info = managed_jobsv1_pb2.ManagedJobInfo(
+                    # The `spot.job_id`, which can be used to identify
+                    # different tasks for the same job
+                    _job_id=job.get('_job_id'),
                     job_id=job.get('job_id'),
                     task_id=job.get('task_id'),
                     job_name=job.get('job_name'),
@@ -466,11 +483,7 @@ class ManagedJobsServiceImpl(managed_jobsv1_pb2_grpc.ManagedJobsServiceServicer
                     end_at=job.get('end_at'),
                     user_yaml=job.get('user_yaml'),
                     entrypoint=job.get('entrypoint'),
-                    metadata={
-                        k: v
-                        for k, v in job.get('metadata', {}).items()
-                        if v is not None
-                    },
+                    metadata=converted_metadata,
                     pool=job.get('pool'),
                     pool_hash=job.get('pool_hash'))
                 jobs_info.append(job_info)

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,74 @@ def daemonize():
     # This process is now fully detached from the original parent and terminal
 
 
-if __name__ == '__main__':
-    daemonize()
+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 +143,40 @@ 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.
+        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 pgid is not None:
+        kill_process_group(pgid)
+    else:
+        kill_process_tree(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
+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()
 
@@ -481,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.
@@ -493,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()
 
@@ -564,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)
@@ -580,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))}')
@@ -666,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:
@@ -880,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/server.py

@@ -7,6 +7,7 @@ import fastapi
 from sky import core as sky_core
 from sky.server.requests import executor
 from sky.server.requests import payloads
+from sky.server.requests import request_names
 from sky.server.requests import requests as requests_lib
 from sky.ssh_node_pools import core as ssh_node_pools_core
 from sky.utils import common_utils
@@ -99,9 +100,9 @@ async def deploy_ssh_node_pool(request: fastapi.Request,
     """Deploy SSH Node Pool using existing ssh_up functionality."""
     try:
         ssh_up_body = payloads.SSHUpBody(infra=pool_name, cleanup=False)
-        executor.schedule_request(
+        await executor.schedule_request_async(
             request_id=request.state.request_id,
-            request_name='ssh_up',
+            request_name=request_names.RequestName.SSH_NODE_POOLS_UP,
             request_body=ssh_up_body,
             func=sky_core.ssh_up,
             schedule_type=requests_lib.ScheduleType.LONG,
@@ -124,9 +125,9 @@ async def deploy_ssh_node_pool_general(
         ssh_up_body: payloads.SSHUpBody) -> Dict[str, str]:
     """Deploys all SSH Node Pools."""
     try:
-        executor.schedule_request(
+        await executor.schedule_request_async(
             request_id=request.state.request_id,
-            request_name='ssh_up',
+            request_name=request_names.RequestName.SSH_NODE_POOLS_UP,
             request_body=ssh_up_body,
             func=sky_core.ssh_up,
             schedule_type=requests_lib.ScheduleType.LONG,
@@ -150,9 +151,9 @@ async def down_ssh_node_pool(request: fastapi.Request,
     """Cleans up a SSH Node Pools."""
     try:
         ssh_up_body = payloads.SSHUpBody(infra=pool_name, cleanup=True)
-        executor.schedule_request(
+        await executor.schedule_request_async(
             request_id=request.state.request_id,
-            request_name='ssh_down',
+            request_name=request_names.RequestName.SSH_NODE_POOLS_DOWN,
             request_body=ssh_up_body,
             func=sky_core.ssh_up,  # Reuse ssh_up function with cleanup=True
             schedule_type=requests_lib.ScheduleType.LONG,
@@ -178,9 +179,9 @@ async def down_ssh_node_pool_general(
     try:
         # Set cleanup=True for down operation
         ssh_up_body.cleanup = True
-        executor.schedule_request(
+        await executor.schedule_request_async(
             request_id=request.state.request_id,
-            request_name='ssh_down',
+            request_name=request_names.RequestName.SSH_NODE_POOLS_DOWN,
             request_body=ssh_up_body,
             func=sky_core.ssh_up,  # Reuse ssh_up function with cleanup=True
             schedule_type=requests_lib.ScheduleType.LONG,