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

sky/utils/controller_utils.py

@@ -23,10 +23,10 @@ from sky.clouds import gcp
 from sky.data import data_utils
 from sky.data import storage as storage_lib
 from sky.jobs import constants as managed_job_constants
-from sky.jobs import state as managed_job_state
 from sky.provision.kubernetes import constants as kubernetes_constants
 from sky.serve import constants as serve_constants
 from sky.serve import serve_state
+from sky.server import config as server_config
 from sky.setup_files import dependencies
 from sky.skylet import constants
 from sky.skylet import log_lib
@@ -72,7 +72,8 @@ class _ControllerSpec:
     """Spec for skypilot controllers."""
     controller_type: str
     name: str
-    cluster_name: str
+    _cluster_name_func: Callable[[], str]
+    _cluster_name_from_server: Optional[str]  # For client-side only
     in_progress_hint: Callable[[bool], str]
     decline_cancel_hint: str
     _decline_down_when_failed_to_fetch_status_hint: str
@@ -93,6 +94,24 @@ class _ControllerSpec:
         return self._check_cluster_name_hint.format(
             cluster_name=self.cluster_name)
 
+    @property
+    def cluster_name(self) -> str:
+        """The cluster name of the controller.
+
+        On the server-side, the cluster name is the actual cluster name,
+        which is read from common.(JOB|SKY_SERVE)_CONTROLLER_NAME.
+
+        On the client-side, the cluster name may not be accurate,
+        as we may not know the exact name, because we are missing
+        the server-side common.SERVER_ID. We have to wait until
+        we get the actual cluster name from the server.
+        """
+        return (self._cluster_name_from_server if self._cluster_name_from_server
+                is not None else self._cluster_name_func())
+
+    def set_cluster_name_from_server(self, cluster_name: str) -> None:
+        self._cluster_name_from_server = cluster_name
+
 
 # TODO: refactor controller class to not be an enum.
 class Controllers(enum.Enum):
@@ -102,7 +121,8 @@ class Controllers(enum.Enum):
     JOBS_CONTROLLER = _ControllerSpec(
         controller_type='jobs',
         name='managed jobs controller',
-        cluster_name=common.JOB_CONTROLLER_NAME,
+        _cluster_name_func=lambda: common.JOB_CONTROLLER_NAME,
+        _cluster_name_from_server=None,
         in_progress_hint=lambda _:
         ('* {job_info}To see all managed jobs: '
          f'{colorama.Style.BRIGHT}sky jobs queue{colorama.Style.RESET_ALL}'),
@@ -133,7 +153,8 @@ class Controllers(enum.Enum):
     SKY_SERVE_CONTROLLER = _ControllerSpec(
         controller_type='serve',
         name='serve controller',
-        cluster_name=common.SKY_SERVE_CONTROLLER_NAME,
+        _cluster_name_func=lambda: common.SKY_SERVE_CONTROLLER_NAME,
+        _cluster_name_from_server=None,
         in_progress_hint=(
             lambda pool:
             (f'* To see detailed pool status: {colorama.Style.BRIGHT}'
@@ -166,7 +187,9 @@ class Controllers(enum.Enum):
         default_autostop_config=serve_constants.CONTROLLER_AUTOSTOP)
 
     @classmethod
-    def from_name(cls, name: Optional[str]) -> Optional['Controllers']:
+    def from_name(cls,
+                  name: Optional[str],
+                  expect_exact_match: bool = True) -> Optional['Controllers']:
         """Check if the cluster name is a controller name.
 
         Returns:
@@ -187,7 +210,11 @@ class Controllers(enum.Enum):
         elif name.startswith(common.JOB_CONTROLLER_PREFIX):
             controller = cls.JOBS_CONTROLLER
             prefix = common.JOB_CONTROLLER_PREFIX
-        if controller is not None and name != controller.value.cluster_name:
+
+        if controller is not None and expect_exact_match:
+            assert name == controller.value.cluster_name, (
+                name, controller.value.cluster_name)
+        elif controller is not None and name != controller.value.cluster_name:
             # The client-side cluster_name is not accurate. Assume that `name`
             # is the actual cluster name, so need to set the controller's
             # cluster name to the input name.
@@ -201,7 +228,7 @@ class Controllers(enum.Enum):
                                                                          prefix)
 
             # Update the cluster name.
-            controller.value.cluster_name = name
+            controller.value.set_cluster_name_from_server(name)
         return controller
 
     @classmethod
@@ -228,10 +255,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.
     """
-    controller = Controllers.from_name(cluster_name)
+    controller = Controllers.from_name(cluster_name, expect_exact_match=False)
     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'),
@@ -400,7 +438,7 @@ def check_cluster_name_not_controller(
     Returns:
       None, if the cluster name is not a controller name.
     """
-    controller = Controllers.from_name(cluster_name)
+    controller = Controllers.from_name(cluster_name, expect_exact_match=False)
     if controller is not None:
         msg = controller.value.check_cluster_name_hint
         if operation_str is not None:
@@ -495,6 +533,9 @@ def shared_controller_vars_to_fill(
         # before popping allowed_contexts. If it is not on Kubernetes,
         # we may be able to use allowed_contexts.
         local_user_config.pop('allowed_contexts', None)
+        # Remove api_server config so that the controller does not try to use
+        # a remote API server.
+        local_user_config.pop('api_server', None)
         with tempfile.NamedTemporaryFile(
                 delete=False,
                 suffix=_LOCAL_SKYPILOT_CONFIG_PATH_SUFFIX) as temp_file:
@@ -528,7 +569,15 @@ def shared_controller_vars_to_fill(
         # with a remote API server.
         constants.USING_REMOTE_API_SERVER_ENV_VAR: str(
             common_utils.get_using_remote_api_server()),
+        constants.IS_SKYPILOT_SERVE_CONTROLLER:
+            ('true'
+             if controller == Controllers.SKY_SERVE_CONTROLLER else 'false'),
     })
+    override_concurrent_launches = os.environ.get(
+        constants.SERVE_OVERRIDE_CONCURRENT_LAUNCHES, None)
+    if override_concurrent_launches is not None:
+        env_vars[constants.SERVE_OVERRIDE_CONCURRENT_LAUNCHES] = str(
+            int(override_concurrent_launches))
     if skypilot_config.loaded():
         # Only set the SKYPILOT_CONFIG env var if the user has a config file.
         env_vars[
@@ -609,15 +658,16 @@ def get_controller_resources(
     controller_resources_to_use: resources.Resources = list(
         controller_resources)[0]
 
-    controller_record = global_user_state.get_cluster_from_name(
+    controller_handle = global_user_state.get_handle_from_cluster_name(
         controller.value.cluster_name)
-    if controller_record is not None:
-        handle = controller_record.get('handle', None)
-        if handle is not None:
+    if controller_handle is not None:
+        if controller_handle is not None:
             # Use the existing resources, but override the autostop config with
             # the one currently specified in the config.
-            controller_resources_to_use = handle.launched_resources.copy(
-                autostop=controller_resources_config_copied.get('autostop'))
+            controller_resources_to_use = (
+                controller_handle.launched_resources.copy(
+                    autostop=controller_resources_config_copied.get('autostop'))
+            )
 
     # If the controller and replicas are from the same cloud (and region/zone),
     # it should provide better connectivity. We will let the controller choose
@@ -714,6 +764,17 @@ def get_controller_resources(
     return result
 
 
+def get_controller_mem_size_gb() -> float:
+    try:
+        with open(os.path.expanduser(constants.CONTROLLER_K8S_MEMORY_FILE),
+                  'r',
+                  encoding='utf-8') as f:
+            return float(f.read())
+    except FileNotFoundError:
+        pass
+    return common_utils.get_mem_size_gb()
+
+
 def _setup_proxy_command_on_controller(
         controller_launched_cloud: 'clouds.Cloud',
         user_config: Dict[str, Any]) -> config_utils.Config:
@@ -1174,77 +1235,175 @@ def maybe_translate_local_file_mounts_and_sync_up(task: 'task_lib.Task',
 
 # ======================= Resources Management Functions =======================
 
-# Based on testing, assume a running job process uses 350MB memory. We use the
-# same estimation for service controller process.
-JOB_MEMORY_MB = 350
-# Monitoring process for service is 1GB. This is based on an old estimation but
-# we keep it here for now.
+# Monitoring process for service is 512MB. This is based on an old
+# estimation but we keep it here for now.
 # TODO(tian): Remeasure this.
-SERVE_MONITORING_MEMORY_MB = 1024
-# The ratio of service controller process to job process. We will treat each
-# service as SERVE_PROC_RATIO job processes.
-SERVE_PROC_RATIO = SERVE_MONITORING_MEMORY_MB / JOB_MEMORY_MB
-# Past 2000 simultaneous jobs, we become unstable.
-# See https://github.com/skypilot-org/skypilot/issues/4649.
-MAX_JOB_LIMIT = 2000
-# Number of ongoing launches launches allowed per CPU, for managed jobs.
-JOB_LAUNCHES_PER_CPU = 4
-# Number of ongoing launches launches allowed per CPU, for services. This is
-# also based on an old estimation, but SKyServe indeed spawn a new process
-# for each launch operation, so it should be slightly more resources demanding
-# than managed jobs.
-SERVE_LAUNCHES_PER_CPU = 2
-# The ratio of service launch to job launch. This is inverted as the parallelism
-# is determined by 1 / LAUNCHES_PER_CPU.
-SERVE_LAUNCH_RATIO = JOB_LAUNCHES_PER_CPU / SERVE_LAUNCHES_PER_CPU
+SERVE_MONITORING_MEMORY_MB = 512
+# The resource consumption ratio of service launch to serve down.
+SERVE_LAUNCH_RATIO = 2.0
 
 # The _RESOURCES_LOCK should be held whenever we are checking the parallelism
 # control or updating the schedule_state of any job or service. Any code that
 # takes this lock must conclude by calling maybe_schedule_next_jobs.
 _RESOURCES_LOCK = '~/.sky/locks/controller_resources.lock'
 
-
-@annotations.lru_cache(scope='global', maxsize=1)
-def get_resources_lock_path() -> str:
-    path = os.path.expanduser(_RESOURCES_LOCK)
-    os.makedirs(os.path.dirname(path), exist_ok=True)
-    return path
+# keep 2GB reserved after the controllers
+MAXIMUM_CONTROLLER_RESERVED_MEMORY_MB = 2048
+
+# NOTE: In the current implementation, we only consider the memory
+# The ratio of resources consumption for managed jobs and pool/serve.
+# This measures pool_resources / jobs_resources. If 2 GB memory is allocated to
+# jobs, then 2 * POOL_JOBS_RESOURCES_RATIO GB memory is allocated to pool/serve.
+POOL_JOBS_RESOURCES_RATIO = 1
+# Number of ongoing launches launches allowed per worker. Can probably be
+# increased a bit to around 16 but keeping it lower to just to be safe
+LAUNCHES_PER_WORKER = 8
+# Number of ongoing launches allowed per service. Can probably be increased
+# a bit as well.
+LAUNCHES_PER_SERVICE = 4
+
+# Based on testing, each worker takes around 200-300MB memory. Keeping it
+# higher to be safe.
+JOB_WORKER_MEMORY_MB = 400
+# this can probably be increased to around 300-400 but keeping it lower to just
+# to be safe
+MAX_JOBS_PER_WORKER = 200
+# Maximum number of controllers that can be running. Hard to handle more than
+# 512 launches at once.
+MAX_CONTROLLERS = 512 // LAUNCHES_PER_WORKER
+# Limit the number of jobs that can be running at once on the entire jobs
+# controller cluster. It's hard to handle cancellation of more than 2000 jobs at
+# once.
+# TODO(cooperc): Once we eliminate static bottlenecks (e.g. sqlite), remove this
+# hardcoded max limit.
+MAX_TOTAL_RUNNING_JOBS = 2000
+
+
+def compute_memory_reserved_for_controllers(
+        reserve_for_controllers: bool, reserve_extra_for_pool: bool) -> float:
+    reserved_memory_mb = 0.0
+    if reserve_for_controllers:
+        reserved_memory_mb = float(MAXIMUM_CONTROLLER_RESERVED_MEMORY_MB)
+        if reserve_extra_for_pool:
+            reserved_memory_mb *= (1. + POOL_JOBS_RESOURCES_RATIO)
+    return reserved_memory_mb
+
+
+def _get_total_usable_memory_mb(pool: bool, consolidation_mode: bool) -> float:
+    controller_reserved = compute_memory_reserved_for_controllers(
+        reserve_for_controllers=True, reserve_extra_for_pool=pool)
+    total_memory_mb = (common_utils.get_mem_size_gb() * 1024 -
+                       controller_reserved)
+    if not consolidation_mode:
+        return total_memory_mb
+    config = server_config.compute_server_config(
+        deploy=True, quiet=True, reserved_memory_mb=controller_reserved)
+    used = 0.0
+    used += ((config.long_worker_config.garanteed_parallelism +
+              config.long_worker_config.burstable_parallelism) *
+             server_config.LONG_WORKER_MEM_GB * 1024)
+    used += ((config.short_worker_config.garanteed_parallelism +
+              config.short_worker_config.burstable_parallelism) *
+             server_config.SHORT_WORKER_MEM_GB * 1024)
+    return total_memory_mb - used
+
+
+def _is_consolidation_mode(pool: bool) -> bool:
+    return skypilot_config.get_nested(
+        ('jobs' if pool else 'serve', 'controller', 'consolidation_mode'),
+        default_value=False)
 
 
 @annotations.lru_cache(scope='request')
-def _get_job_parallelism() -> int:
-    job_memory = JOB_MEMORY_MB * 1024 * 1024
-    job_limit = min(psutil.virtual_memory().total // job_memory, MAX_JOB_LIMIT)
-    return max(job_limit, 1)
+def _get_parallelism(pool: bool, raw_resource_per_unit: float) -> int:
+    """Returns the number of jobs controllers / services that should be running.
 
+    This is the number of controllers / services that should be running
+    to maximize resource utilization.
 
-@annotations.lru_cache(scope='request')
-def _get_launch_parallelism() -> int:
-    cpus = os.cpu_count()
-    return cpus * JOB_LAUNCHES_PER_CPU if cpus is not None else 1
+    In consolidation mode, we use the existing API server so our resource
+    requirements are just for the job controllers / services. We try taking
+    up as much memory as possible left over from the API server.
 
+    In non-consolidation mode, we have to take into account the memory of the
+    API server workers. We limit to only 8 launches per worker, so our logic is
+    each controller will take CONTROLLER_MEMORY_MB + 8 * WORKER_MEMORY_MB. We
+    leave some leftover room for ssh codegen and ray status overhead.
+    """
+    consolidation_mode = _is_consolidation_mode(pool)
 
-def can_provision() -> bool:
-    # We always prioritize terminating over provisioning, to save the cost on
-    # idle resources.
-    if serve_state.total_number_scheduled_to_terminate_replicas() > 0:
-        return False
-    return can_terminate()
+    total_memory_mb = _get_total_usable_memory_mb(pool, consolidation_mode)
+
+    # In consolidation mode, we assume the API server is running in deployment
+    # mode, hence resource management (i.e. how many requests are allowed) is
+    # done by the API server.
+    resource_per_unit_worker = 0.
+    # Otherwise, it runs a local API server on the jobs/serve controller.
+    # We need to do the resource management ourselves.
+    if not consolidation_mode:
+        launches_per_worker = (LAUNCHES_PER_WORKER
+                               if pool else LAUNCHES_PER_SERVICE)
+        resource_per_unit_worker = (launches_per_worker *
+                                    server_config.LONG_WORKER_MEM_GB * 1024)
 
+    # If running pool on jobs controller, we need to account for the resources
+    # consumed by the jobs.
+    ratio = (1. + POOL_JOBS_RESOURCES_RATIO) if pool else 1.
+    resource_per_unit = ratio * (raw_resource_per_unit +
+                                 resource_per_unit_worker)
 
-def can_start_new_process() -> bool:
-    num_procs = (serve_state.get_num_services() * SERVE_PROC_RATIO +
-                 managed_job_state.get_num_alive_jobs())
-    return num_procs < _get_job_parallelism()
+    return max(int(total_memory_mb / resource_per_unit), 1)
 
 
-# We limit the number of terminating replicas to the number of CPUs. This is
-# just a temporary solution to avoid overwhelming the controller. After one job
-# controller PR, we should use API server to handle resources management.
-def can_terminate() -> bool:
+def get_number_of_jobs_controllers() -> int:
+    return min(
+        MAX_CONTROLLERS,
+        _get_parallelism(pool=True, raw_resource_per_unit=JOB_WORKER_MEMORY_MB))
+
+
+@annotations.lru_cache(scope='global', maxsize=1)
+def get_resources_lock_path() -> str:
+    path = os.path.expanduser(_RESOURCES_LOCK)
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    return path
+
+
+def _get_number_of_services(pool: bool) -> int:
+    return _get_parallelism(pool=pool,
+                            raw_resource_per_unit=SERVE_MONITORING_MEMORY_MB *
+                            POOL_JOBS_RESOURCES_RATIO)
+
+
+@annotations.lru_cache(scope='request')
+def _get_request_parallelism(pool: bool) -> int:
+    # NOTE(dev): One smoke test depends on this value.
+    # tests/smoke_tests/test_sky_serve.py::test_skyserve_new_autoscaler_update
+    # assumes 4 concurrent launches.
+    override_concurrent_launches = os.environ.get(
+        constants.SERVE_OVERRIDE_CONCURRENT_LAUNCHES, None)
+    if override_concurrent_launches is not None and not pool:
+        return int(override_concurrent_launches)
+    # Limitation per service x number of services
+    launches_per_worker = (LAUNCHES_PER_WORKER
+                           if pool else LAUNCHES_PER_SERVICE)
+    return (launches_per_worker * POOL_JOBS_RESOURCES_RATIO *
+            _get_number_of_services(pool))
+
+
+def can_provision(pool: bool) -> bool:
+    # TODO(tian): probe API server to see if there is any pending provision
+    # requests.
+    return can_terminate(pool)
+
+
+def can_start_new_process(pool: bool) -> bool:
+    return serve_state.get_num_services() < _get_number_of_services(pool)
+
+
+def can_terminate(pool: bool) -> bool:
+    # TODO(tian): probe API server to see if there is any pending terminate
+    # requests.
     num_terminating = (
-        serve_state.total_number_provisioning_replicas() * SERVE_LAUNCH_RATIO +
-        # Each terminate process will take roughly the same CPUs as job launch.
-        serve_state.total_number_terminating_replicas() +
-        managed_job_state.get_num_launching_jobs())
-    return num_terminating < _get_launch_parallelism()
+        serve_state.total_number_provisioning_replicas() +
+        serve_state.total_number_terminating_replicas() / SERVE_LAUNCH_RATIO)
+    return num_terminating < _get_request_parallelism(pool)

sky/utils/db/db_utils.py

@@ -7,15 +7,17 @@ import pathlib
 import sqlite3
 import threading
 import typing
-from typing import Any, Callable, Dict, Iterable, Optional
+from typing import Any, Callable, Dict, Iterable, Literal, Optional, Union
 
 import aiosqlite
 import aiosqlite.context
 import sqlalchemy
 from sqlalchemy import exc as sqlalchemy_exc
+from sqlalchemy.ext import asyncio as sqlalchemy_async
 
 from sky import sky_logging
 from sky.skylet import constants
+from sky.skylet import runtime_utils
 
 logger = sky_logging.init_logger(__name__)
 if typing.TYPE_CHECKING:
@@ -184,7 +186,7 @@ def add_column_to_table_sqlalchemy(
             pass
         else:
             raise
-    #postgressql
+    #postgresql
     except sqlalchemy_exc.ProgrammingError as e:
         if 'already exists' in str(e):
             pass
@@ -200,6 +202,7 @@ def add_column_to_table_alembic(
     server_default: Optional[str] = None,
     copy_from: Optional[str] = None,
     value_to_replace_existing_entries: Optional[Any] = None,
+    index: Optional[bool] = None,
 ):
     """Add a column to a table using Alembic operations.
 
@@ -214,6 +217,8 @@ def add_column_to_table_alembic(
         copy_from: Column name to copy values from (for existing rows)
         value_to_replace_existing_entries: Default value for existing NULL
             entries
+        index: If True, create an index on this column. If None, no index
+            is created.
     """
     from alembic import op  # pylint: disable=import-outside-toplevel
 
@@ -221,7 +226,8 @@ def add_column_to_table_alembic(
         # Create the column with server_default if provided
         column = sqlalchemy.Column(column_name,
                                    column_type,
-                                   server_default=server_default)
+                                   server_default=server_default,
+                                   index=index)
         op.add_column(table_name, column)
 
         # Handle data migration
@@ -353,6 +359,27 @@ class SQLiteConn(threading.local):
         conn = await self._get_async_conn()
         return await conn.execute_fetchall(sql, parameters)
 
+    async def execute_get_returning_value_async(
+            self,
+            sql: str,
+            parameters: Optional[Iterable[Any]] = None
+    ) -> Optional[sqlite3.Row]:
+        conn = await self._get_async_conn()
+
+        if parameters is None:
+            parameters = []
+
+        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
+
+        # pylint: disable=protected-access
+        return await conn._execute(exec_and_get_returning_value, sql,
+                                   parameters)
+
     async def close(self):
         if self._async_conn is not None:
             await self._async_conn.close()
@@ -375,32 +402,82 @@ def get_max_connections():
     return _max_connections
 
 
-def get_engine(db_name: str):
+@typing.overload
+def get_engine(
+        db_name: Optional[str],
+        async_engine: Literal[False] = False) -> sqlalchemy.engine.Engine:
+    ...
+
+
+@typing.overload
+def get_engine(db_name: Optional[str],
+               async_engine: Literal[True]) -> sqlalchemy_async.AsyncEngine:
+    ...
+
+
+def get_engine(
+    db_name: Optional[str],
+    async_engine: bool = False
+) -> Union[sqlalchemy.engine.Engine, sqlalchemy_async.AsyncEngine]:
+    """Get the engine for the given database name.
+
+    Args:
+        db_name: The name of the database. ONLY used for SQLite. On Postgres,
+        we use a single database, which we get from the connection string.
+        async_engine: Whether to return an async engine.
+    """
     conn_string = None
     if os.environ.get(constants.ENV_VAR_IS_SKYPILOT_SERVER) is not None:
         conn_string = os.environ.get(constants.ENV_VAR_DB_CONNECTION_URI)
     if conn_string:
+        if async_engine:
+            conn_string = conn_string.replace('postgresql://',
+                                              'postgresql+asyncpg://')
         with _db_creation_lock:
+            # We use the same cache for both sync and async engines
+            # because we change the conn_string in the async case,
+            # so they would not overlap.
             if conn_string not in _postgres_engine_cache:
+                engine_type = 'sync' if not async_engine else 'async'
+                logger.debug(
+                    f'Creating a new postgres {engine_type} engine with '
+                    f'maximum {_max_connections} connections')
                 if _max_connections == 0:
-                    _postgres_engine_cache[conn_string] = (
-                        sqlalchemy.create_engine(
-                            conn_string, poolclass=sqlalchemy.pool.NullPool))
-                elif _max_connections == 1:
-                    _postgres_engine_cache[conn_string] = (
-                        sqlalchemy.create_engine(
-                            conn_string, poolclass=sqlalchemy.pool.StaticPool))
+                    kw_args = {'poolclass': sqlalchemy.NullPool}
+                    if async_engine:
+                        _postgres_engine_cache[conn_string] = (
+                            sqlalchemy_async.create_async_engine(
+                                conn_string, **kw_args))
+                    else:
+                        _postgres_engine_cache[conn_string] = (
+                            sqlalchemy.create_engine(conn_string, **kw_args))
                 else:
-                    _postgres_engine_cache[conn_string] = (
-                        sqlalchemy.create_engine(
-                            conn_string,
-                            poolclass=sqlalchemy.pool.QueuePool,
-                            size=_max_connections,
-                            max_overflow=0))
+                    kw_args = {
+                        'pool_size': _max_connections,
+                        'max_overflow': max(0, 5 - _max_connections),
+                        'pool_pre_ping': True,
+                        'pool_recycle': 1800
+                    }
+                    if async_engine:
+                        kw_args[
+                            'poolclass'] = sqlalchemy.pool.AsyncAdaptedQueuePool
+                        _postgres_engine_cache[conn_string] = (
+                            sqlalchemy_async.create_async_engine(
+                                conn_string, **kw_args))
+                    else:
+                        kw_args['poolclass'] = sqlalchemy.pool.QueuePool
+                        _postgres_engine_cache[conn_string] = (
+                            sqlalchemy.create_engine(conn_string, **kw_args))
             engine = _postgres_engine_cache[conn_string]
     else:
-        db_path = os.path.expanduser(f'~/.sky/{db_name}.db')
+        assert db_name is not None, 'db_name must be provided for SQLite'
+        db_path = runtime_utils.get_runtime_dir_path(f'.sky/{db_name}.db')
         pathlib.Path(db_path).parents[0].mkdir(parents=True, exist_ok=True)
+        if async_engine:
+            # This is an AsyncEngine, instead of a (normal, synchronous) Engine,
+            # so we should not put it in the cache. Instead, just return.
+            return sqlalchemy_async.create_async_engine(
+                'sqlite+aiosqlite:///' + db_path, connect_args={'timeout': 30})
         if db_path not in _sqlite_engine_cache:
             _sqlite_engine_cache[db_path] = sqlalchemy.create_engine(
                 'sqlite:///' + db_path)