From 977ac6a7f10fcd75ab914b94873d6dc540500cbc Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Mon, 14 Sep 2026 15:17:52 +0500 Subject: [PATCH 1/8] Document how to work with DB connections --- AGENTS.md | 4 +- contributing/DATABASE.md | 84 ++++++++++++++++++++++++++++++++++++++++ contributing/LOCKING.md | 28 +++++++++++++- 3 files changed, 113 insertions(+), 3 deletions(-) create mode 100644 contributing/DATABASE.md diff --git a/AGENTS.md b/AGENTS.md index 6b7aec13cc..234464fbd1 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,7 +1,7 @@ # Repository Guidelines Before touching a subsystem, read the relevant notes in `contributing/`: `ARCHITECTURE.md`, -`PIPELINES.md`, `LOCKING.md`, `MIGRATIONS.md`, `RUNS-AND-JOBS.md`, `AUTOSCALING.md`, +`PIPELINES.md`, `LOCKING.md`, `DATABASE.md`, `MIGRATIONS.md`, `RUNS-AND-JOBS.md`, `AUTOSCALING.md`, `BACKENDS.md`, `GPUHUNT.md`, `PROXY.md`, `RUNNER-AND-SHIM.md`, `FRONTEND.md`, `DOCS.md`, `DEVELOPMENT.md`, `RELEASE.md`. @@ -29,7 +29,7 @@ Before touching a subsystem, read the relevant notes in `contributing/`: `ARCHIT - Prefer pydantic-style models in `core/models`. - Document attributes when the note adds behavior, compatibility, or semantic context that is not obvious from the name and type. Use attribute docstrings without leading newline. - Tests use `test_*.py` modules and `test_*` functions; fixtures live near usage. -- Never make network calls inside a DB session or transaction. Fetch what you need before opening the session, or commit and close it before the call. +- Never make network calls inside a DB session or transaction. Fetch what you need before opening the session, or commit and close it before the call. See `contributing/DATABASE.md`. - Don't use function-level (inner) imports to break circular imports. Inject the dependency or move the shared code to a lower-level module instead. - Never edit a migration that has already been applied or released; add a new migration instead. - Derive paths under `SERVER_DIR_PATH` on access (a `get_*` function), not as module-level constants, so that patching `settings.SERVER_DIR_PATH` redirects all of them. Tests rely on this to keep server state out of the real `~/.dstack`. diff --git a/contributing/DATABASE.md b/contributing/DATABASE.md new file mode 100644 index 0000000000..484fb79b9c --- /dev/null +++ b/contributing/DATABASE.md @@ -0,0 +1,84 @@ +# Database + +This document describes the `dstack` server rules of working with the database. Separately covered: + +* DB migrations in `MIGRATIONS.md` +* Resource locking in `LOCKING.md` +* Pipeline-based background processing in `PIPELINES.md`. + +## Connections + +DB connections are a scarce resource. For example, a minimal RDS Postgres instance comes with ~80 connections. Increasing the number of connections requires +scaling the DB instance and is expensive. This section describes how the `dstack` server should use DB connections to avoid running out of them. + +### Connection budget + +The server creates one SQLAlchemy async engine per process with an `AsyncAdaptedQueuePool` of `DSTACK_DB_POOL_SIZE` (default 20) persistent connections plus `DSTACK_DB_MAX_OVERFLOW` (default 20) overflow connections, i.e. up to 40 connections per server process. When all of them are checked out, the next checkout waits `pool_timeout` (30 seconds) and then fails with: + +``` +sqlalchemy.exc.TimeoutError: QueuePool limit of size 20 overflow 20 reached, connection timed out, timeout 30.00 +``` + +Far more coroutines compete for the slots than there are slots. This works because each of them holds a connection for a short time. + +A session does not hold a connection until it executes its first statement. From that moment until the session commits, rolls back, or closes, the connection is checked out and unavailable to anyone else. The invariant that follows: + +**Hold a connection only while doing DB work. Do not keep a session open across when doing heavy work.** + +On SQLite the same rule applies for a different reason: a long write transaction blocks every other writer for up to `busy_timeout` (30 seconds). + +### Awaits that must not run inside a session + +Anything whose completion depends on something other than the database: + +* SSH: opening tunnels, `aexec()`, `acheck()`, any subprocess. +* HTTP to runners, shims, gateways, or external services. +* Cloud SDK calls, including those wrapped in `run_async()`. The thread pool executor is shared, so a saturated executor makes even a fast call wait. +* Retry loops with sleeps, and waiting on asyncio locks or events that other tasks control. +* Streaming a response to a client. + +Fetch what you need before opening the session, or commit and close the session before the call and open a new one to apply the results: + +```python +async with get_session_ctx() as session: + # Load and lock what processing needs. + context = await _load_context(session, item) + +# No connection is held here. +result = await _do_network_work(context) + +async with get_session_ctx() as session: + # Apply results with a guarded update. + await _apply_result(session, item, result) +``` + +### Session lifetime by context + +**Pipelines** + +Pipeline workers follow load, process, apply with a short session for the load and apply phases and no session during processing. See `PIPELINES.md`. + +**API handlers** + +The request session is created before the handler runs. The authentication dependency executes a SELECT on it, so a connection is checked out before the handler starts and stays checked out until the handler returns and the session commits. + +This is acceptable because API handlers are assumed to be quick: read or write a few rows and return. A handler that needs to do slow work must not do it while holding the request session's connection. Either: + +* Commit the session before the slow work. Committing returns the connection to the pool. The next statement on the session checks out a new connection and starts a new transaction, so any locks taken earlier are released, and the handler must not rely on them across the gap. +* Make the API async and move the slow work out of the handler. + +The rule is not applied currently applied to API handlers strictly: a handler may make one or a few network calls while holding the connection. + +**Startup** + +Migrations and server initialization may hold a connection and a session-level advisory lock for long: they run once per process, before anything else competes for connections. Recurring cross-replica coordination should use the pipeline lock columns instead. + +### Lock waits + +A coroutine blocked on a DB-level lock holds its connection while it waits. This applies to `SELECT ... FOR UPDATE` and to `pg_advisory_lock` alike: the waiter occupies a pool slot on its replica until the holder releases the lock. A lock held for a long time thus costs one slot on the holder plus one per waiter, and a hung holder pins all of them. + +Keep locked regions short and bounded. See `LOCKING.md` for how to take advisory locks correctly. + +### Settings + +`DSTACK_DB_POOL_SIZE` and `DSTACK_DB_MAX_OVERFLOW` size the pool per process. Raising them lets more coroutines hold connections at once, which hides a long hold rather than fixing it, and each replica needs `pool_size + max_overflow` connections from the Postgres `max_connections` budget. diff --git a/contributing/LOCKING.md b/contributing/LOCKING.md index c89ad526a4..604b1a211b 100644 --- a/contributing/LOCKING.md +++ b/contributing/LOCKING.md @@ -29,7 +29,7 @@ Locksets are an optimization. One can think of them as per-resource-id locks tha Postgres resource locking is implemented via standard SELECT FOR UPDATE. SQLAlchemy provides `.with_for_update()` that has no effect if SELECT FOR UPDATE is not supported as in SQLite. -There are few places that rely on advisory locks as when generating unique resource names. +There are few places that rely on advisory locks as when generating unique resource names or serializing server initialization across replicas. See **Advisory locks** below. ## Working with locks @@ -112,6 +112,32 @@ Note that: * This pattern works assuming that Postgres is using default isolation level Read Committed. By the time a transaction acquires the advisory lock, all other transactions that can take the name have committed, so their changes can be seen and a unique name is taken. * SQLite needs a commit before selecting taken names due to Snapshot Isolation as noted above. +**Advisory locks** + +Postgres has two kinds of advisory locks: + +* `pg_advisory_xact_lock` is released when the transaction ends. The unique names pattern above uses it. +* `pg_advisory_lock` is bound to the connection. It survives commit and rollback and is released only by `pg_advisory_unlock` on the same connection or when the connection closes. `advisory_lock_ctx()` in `services/locking.py` wraps this kind for work that must span several transactions, such as migrations or server initialization. + +Either use `pg_advisory_xact_lock` within a single transaction, or use `advisory_lock_ctx()` and follow these rules: + +* Bind it to an `AsyncConnection` from `engine.connect()`, not to an `AsyncSession`. If the session commits inside the block, its next statement may run on a different pooled connection, and `pg_advisory_unlock` goes to a connection that never held the lock. Postgres only returns `false` with a warning in this case, so the failure is silent: the lock stays on an idle pooled connection until the process exits, and every replica blocks forever on its next acquire. See https://github.com/dstackai/dstack/issues/3881 for an example. +* Keep the locked block short and bounded. Every waiter is blocked inside `pg_advisory_lock` holding a DB connection of its own, so a long or hung holder pins one connection per waiter. See `DATABASE.md`. + +```python +async with get_db().engine.connect() as connection: + async with advisory_lock_ctx( + bind=connection, + dialect_name=get_db().dialect_name, + resource="server_init", + ): + async with get_session_ctx() as session: + # The session may commit freely: the lock lives on `connection`. + ... +``` + +A released connection goes back to the pool, so a lock that failed to release stays there too. `_release_advisory_lock()` tolerates failures because the common one is an invalidated connection, in which case Postgres has already dropped the lock. A release that fails on a live connection strands the lock. + **Use `AsyncExitStack`** In-memory locking typically requires taking lock for long (until commit). From 92a6cfa0526c620ddc2fc1661901623ef98fcc50 Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Mon, 14 Sep 2026 15:45:38 +0500 Subject: [PATCH 2/8] Use dedicated connection for server_init lock --- src/dstack/_internal/server/app.py | 51 ++++++++++++++++-------------- 1 file changed, 28 insertions(+), 23 deletions(-) diff --git a/src/dstack/_internal/server/app.py b/src/dstack/_internal/server/app.py index da5877a277..0b06dd6b20 100644 --- a/src/dstack/_internal/server/app.py +++ b/src/dstack/_internal/server/app.py @@ -130,33 +130,38 @@ async def lifespan(app: FastAPI): server_config_loaded = server_config_manager.load_config() # Encryption has to be configured before working with users and projects await server_config_manager.apply_encryption() - async with get_session_ctx() as session: + async with get_db().engine.connect() as lock_connection: + # The lock is bound to a dedicated connection because the session commits inside the block, + # which would move a session-bound lock release to a different pooled connection. async with advisory_lock_ctx( - bind=session, + bind=lock_connection, dialect_name=get_db().dialect_name, resource="server_init", ): - admin, _ = await get_or_create_admin_user(session=session) - await get_or_create_default_project( - session=session, - user=admin, - ) - if server_config_manager is not None: - server_config_file_path = get_server_config_file_path() - server_config_dir = _get_server_config_dir() - if not server_config_loaded: - logger.info("Initializing the default configuration...", {"show_path": False}) - await server_config_manager.init_config(session=session) - logger.info( - f"Initialized the default configuration at [link=file://{server_config_file_path}]{server_config_dir}[/link]", - {"show_path": False}, - ) - else: - logger.info( - f"Applying [link=file://{server_config_file_path}]{server_config_dir}[/link]...", - {"show_path": False}, - ) - await server_config_manager.apply_config(session=session, owner=admin) + async with get_session_ctx() as session: + admin, _ = await get_or_create_admin_user(session=session) + await get_or_create_default_project( + session=session, + user=admin, + ) + if server_config_manager is not None: + server_config_file_path = get_server_config_file_path() + server_config_dir = _get_server_config_dir() + if not server_config_loaded: + logger.info( + "Initializing the default configuration...", {"show_path": False} + ) + await server_config_manager.init_config(session=session) + logger.info( + f"Initialized the default configuration at [link=file://{server_config_file_path}]{server_config_dir}[/link]", + {"show_path": False}, + ) + else: + logger.info( + f"Applying [link=file://{server_config_file_path}]{server_config_dir}[/link]...", + {"show_path": False}, + ) + await server_config_manager.apply_config(session=session, owner=admin) update_default_project( project_name=DEFAULT_PROJECT_NAME, From 8dd599e8e843690a9e19fbad8f80e6c442340fc3 Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Mon, 14 Sep 2026 16:37:46 +0500 Subject: [PATCH 3/8] Set command_timeout on Postgres --- contributing/DATABASE.md | 2 ++ mkdocs/docs/reference/env.md | 1 + src/dstack/_internal/server/db.py | 10 +++++++++- src/dstack/_internal/server/settings.py | 4 ++++ 4 files changed, 16 insertions(+), 1 deletion(-) diff --git a/contributing/DATABASE.md b/contributing/DATABASE.md index 484fb79b9c..f70dffa084 100644 --- a/contributing/DATABASE.md +++ b/contributing/DATABASE.md @@ -81,4 +81,6 @@ Keep locked regions short and bounded. See `LOCKING.md` for how to take advisory ### Settings +`DSTACK_DB_COMMAND_TIMEOUT` (default 300 seconds) bounds every operation on a Postgres connection, so a coroutine waiting on a connection whose socket died silently gets an error and releases the slot instead of holding it forever. It applies to migrations too, since they run on the same engine, so it must leave room for index builds and backfills. + `DSTACK_DB_POOL_SIZE` and `DSTACK_DB_MAX_OVERFLOW` size the pool per process. Raising them lets more coroutines hold connections at once, which hides a long hold rather than fixing it, and each replica needs `pool_size + max_overflow` connections from the Postgres `max_connections` budget. diff --git a/mkdocs/docs/reference/env.md b/mkdocs/docs/reference/env.md index 5e4a80ee58..f00b0dd96f 100644 --- a/mkdocs/docs/reference/env.md +++ b/mkdocs/docs/reference/env.md @@ -141,6 +141,7 @@ For more details on the options below, refer to the [server deployment](../guide - `DSTACK_SERVER_GCS_BUCKET`{ #DSTACK_SERVER_GCS_BUCKET } - The bucket that repo diffs will be uploaded to if set. If unset, diffs are uploaded to the database. - `DSTACK_DB_POOL_SIZE`{ #DSTACK_DB_POOL_SIZE } - The client DB connections pool size. Defaults to `20`, - `DSTACK_DB_MAX_OVERFLOW`{ #DSTACK_DB_MAX_OVERFLOW } - The client DB connections pool allowed overflow. Defaults to `20`. +- `DSTACK_DB_COMMAND_TIMEOUT`{ #DSTACK_DB_COMMAND_TIMEOUT } - The timeout for a single DB operation in seconds, Postgres only. Set to `0` to disable. Defaults to `300`. - `DSTACK_SERVER_BACKGROUND_PROCESSING_DISABLED`{ #DSTACK_SERVER_BACKGROUND_PROCESSING_DISABLED } - Disables background processing if set to any value. Useful to run only web frontend and API server. - `DSTACK_SERVER_MAX_PROBES_PER_JOB`{ #DSTACK_SERVER_MAX_PROBES_PER_JOB } - Maximum number of probes allowed in a run configuration. Validated at apply time. - `DSTACK_SERVER_MAX_PROBE_TIMEOUT`{ #DSTACK_SERVER_MAX_PROBE_TIMEOUT } - Maximum allowed timeout for a probe. Validated at apply time. diff --git a/src/dstack/_internal/server/db.py b/src/dstack/_internal/server/db.py index 3136b4e6b9..c70a345a22 100644 --- a/src/dstack/_internal/server/db.py +++ b/src/dstack/_internal/server/db.py @@ -2,7 +2,7 @@ from typing import Optional from alembic import command, config -from sqlalchemy import AsyncAdaptedQueuePool, event +from sqlalchemy import AsyncAdaptedQueuePool, event, make_url from sqlalchemy.engine.interfaces import DBAPIConnection from sqlalchemy.ext.asyncio import ( AsyncEngine, @@ -28,6 +28,7 @@ def __init__(self, url: str, engine: Optional[AsyncEngine] = None): poolclass=AsyncAdaptedQueuePool, pool_size=settings.DB_POOL_SIZE, max_overflow=settings.DB_MAX_OVERFLOW, + connect_args=self._get_connect_args(self.url), ) self.session_maker = async_sessionmaker( bind=self.engine, # type: ignore[assignment] @@ -55,6 +56,13 @@ def dialect_name(self) -> str: def get_session(self) -> AsyncSession: return self.session_maker() + def _get_connect_args(self, url: str) -> dict: + if make_url(url).get_backend_name() == "postgresql": + # TODO: Consider setting "command_timeout" high for migrations + # and low for queries – requires a separate Database instance for migrations. + return {"command_timeout": settings.DB_COMMAND_TIMEOUT} + return {} + def get_new_db() -> Database: """ diff --git a/src/dstack/_internal/server/settings.py b/src/dstack/_internal/server/settings.py index 3605725357..e7fe97124b 100644 --- a/src/dstack/_internal/server/settings.py +++ b/src/dstack/_internal/server/settings.py @@ -5,6 +5,7 @@ import os from enum import Enum from pathlib import Path +from typing import Optional from dstack._internal.server.utils.settings import parse_hostname_port from dstack._internal.utils.env import environ @@ -67,6 +68,9 @@ def get_database_url() -> str: # or increase client pool size to support more concurrent requests. DB_POOL_SIZE = int(os.getenv("DSTACK_DB_POOL_SIZE", 20)) DB_MAX_OVERFLOW = int(os.getenv("DSTACK_DB_MAX_OVERFLOW", 20)) +DB_COMMAND_TIMEOUT: Optional[float] = float(os.getenv("DSTACK_DB_COMMAND_TIMEOUT", 300)) or None +"""Bounds every operation on a Postgres connection so that a connection whose +socket died silently is released back to the pool instead of being held forever. 0 disables the timeout.""" SERVER_BACKGROUND_PROCESSING_DISABLED = ( os.getenv("DSTACK_SERVER_BACKGROUND_PROCESSING_DISABLED") is not None From be81404ef22f42a4ee09d5681e546e0e2e4d0dce Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Tue, 15 Sep 2026 10:57:48 +0500 Subject: [PATCH 4/8] Fetch gateway stats outside DB session in runs pipeline --- .../pipeline_tasks/runs/__init__.py | 39 +++++++++---------- .../server/services/gateways/__init__.py | 1 - 2 files changed, 18 insertions(+), 22 deletions(-) diff --git a/src/dstack/_internal/server/background/pipeline_tasks/runs/__init__.py b/src/dstack/_internal/server/background/pipeline_tasks/runs/__init__.py index 20461f0814..a40b003db4 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/runs/__init__.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/runs/__init__.py @@ -18,7 +18,8 @@ import dstack._internal.server.background.pipeline_tasks.runs.active as active import dstack._internal.server.background.pipeline_tasks.runs.pending as pending import dstack._internal.server.background.pipeline_tasks.runs.terminating as terminating -from dstack._internal.core.models.runs import JobStatus, RunStatus +from dstack._internal.core.models.runs import JobStatus, RunSpec, RunStatus +from dstack._internal.proxy.gateway.schemas.stats import PerWindowStats from dstack._internal.server.background.pipeline_tasks.base import ( Fetcher, Heartbeater, @@ -300,6 +301,7 @@ async def _process_pending_item(item: RunPipelineItem) -> None: if context is None: return + context.gateway_stats = await _load_gateway_stats(context.run_model, context.run_spec) result = await pending.process_pending_run(context) if result is None: await _apply_noop_result( @@ -330,21 +332,11 @@ async def _load_pending_context( return None secrets = await get_project_secrets_mapping(session=session, project=run_model.project) run_spec = get_run_spec(run_model) - - gateway_stats = None - if run_spec.configuration.type == "service" and run_model.gateway is not None: - gateway_stats = await get_combined_gateway_stats( - get_gateway_replica_models(run_model.gateway), - run_model.project.name, - run_model.run_name, - ) - return pending.PendingContext( run_model=run_model, run_spec=run_spec, secrets=secrets, locked_job_ids=locked_job_ids, - gateway_stats=gateway_stats, ) @@ -506,6 +498,7 @@ async def _process_active_item(item: RunPipelineItem) -> None: return context = load_result + context.gateway_stats = await _load_gateway_stats(context.run_model, context.run_spec) result = await active.process_active_run(context) await _apply_active_result(item=item, context=context, result=result) @@ -532,21 +525,25 @@ async def _load_active_context( return None secrets = await get_project_secrets_mapping(session=session, project=run_model.project) run_spec = get_run_spec(run_model) - - gateway_stats = None - if run_spec.configuration.type == "service" and run_model.gateway is not None: - gateway_stats = await get_combined_gateway_stats( - get_gateway_replica_models(run_model.gateway), - run_model.project.name, - run_model.run_name, - ) - return active.ActiveContext( run_model=run_model, run_spec=run_spec, secrets=secrets, locked_job_ids=locked_job_ids, - gateway_stats=gateway_stats, + ) + + +async def _load_gateway_stats(run_model: RunModel, run_spec: RunSpec) -> Optional[PerWindowStats]: + """ + Fetches service stats from the gateway replicas. + Talks to gateways over SSH, so it must run outside DB sessions. + """ + if run_spec.configuration.type != "service" or run_model.gateway is None: + return None + return await get_combined_gateway_stats( + get_gateway_replica_models(run_model.gateway), + run_model.project.name, + run_model.run_name, ) diff --git a/src/dstack/_internal/server/services/gateways/__init__.py b/src/dstack/_internal/server/services/gateways/__init__.py index 7b0705faee..e903bbfe07 100644 --- a/src/dstack/_internal/server/services/gateways/__init__.py +++ b/src/dstack/_internal/server/services/gateways/__init__.py @@ -677,7 +677,6 @@ async def generate_gateway_name(session: AsyncSession, project: ProjectModel) -> return name -# TODO: Connect to gateway outside session async def get_or_add_gateway_connections( gateway_replicas: Sequence[GatewayReplicaModel], ) -> List[GatewayConnection]: From 62781e12f3e44dc88bf86b4dbd87b9beb748dedc Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Tue, 15 Sep 2026 13:20:45 +0500 Subject: [PATCH 5/8] Poll migrations lock to avoid deadlock with concurrent index builds A replica blocked in pg_advisory_lock() holds a snapshot that another replica's CREATE INDEX CONCURRENTLY migration waits for, so both hang until Postgres kills the waiter. --- src/dstack/_internal/server/db.py | 39 ++++++++++++++++++++++++------- 1 file changed, 31 insertions(+), 8 deletions(-) diff --git a/src/dstack/_internal/server/db.py b/src/dstack/_internal/server/db.py index c70a345a22..ce9fc82067 100644 --- a/src/dstack/_internal/server/db.py +++ b/src/dstack/_internal/server/db.py @@ -1,3 +1,4 @@ +import asyncio from contextlib import asynccontextmanager from typing import Optional @@ -13,7 +14,7 @@ from sqlalchemy.pool import ConnectionPoolEntry from dstack._internal.server import settings -from dstack._internal.server.services.locking import advisory_lock_ctx +from dstack._internal.server.services.locking import try_advisory_lock_ctx class Database: @@ -88,15 +89,37 @@ def override_db(new_db: Database): _db = new_db +_MIGRATIONS_LOCK_POLL_INTERVAL = 1 +_MIGRATIONS_LOCK_MAX_ATTEMPTS = 600 + + async def migrate(): db = get_db() - async with db.engine.connect() as connection: - async with advisory_lock_ctx( - bind=connection, - dialect_name=db.dialect_name, - resource="migrations", - ): - await connection.run_sync(_run_alembic_upgrade) + # The lock is polled instead of awaited in pg_advisory_lock() to avoid this: + # migrations waiting for older snapshots to finish (e.g. CREATE INDEX CONCURRENTLY) + # wait for the blocked replicas waiting for the "migrations" lock, deadlocking both. + # The lock connection runs in autocommit mode because an idle + # transaction between attempts would keep its snapshot and hang both replicas the same way. + # Migrations run on a separate, transactional connection. + async with db.engine.connect() as lock_connection: + await lock_connection.execution_options(isolation_level="AUTOCOMMIT") + for _ in range(_MIGRATIONS_LOCK_MAX_ATTEMPTS): + async with try_advisory_lock_ctx( + bind=lock_connection, + dialect_name=db.dialect_name, + resource="migrations", + ) as locked: + if locked: + async with db.engine.connect() as connection: + await connection.run_sync(_run_alembic_upgrade) + return + await asyncio.sleep(_MIGRATIONS_LOCK_POLL_INTERVAL) + raise TimeoutError( + "Timed out waiting for the migrations lock." + " Another server replica may be running long migrations, or a replica that lost" + " connectivity may still hold it: check pg_locks for the advisory lock" + " and terminate the holder's backend if it is gone." + ) async def get_session(): From d2ee3451e3986236e9490043a94a95c7e68b9d5c Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Tue, 15 Sep 2026 14:08:27 +0500 Subject: [PATCH 6/8] Run migrations and server init on the connection holding the lock Holding a session-level advisory lock on a separate connection lets Postgres release it, e.g. via idle_in_transaction_session_timeout, while the work continues on another connection. Migrations now run on the lock connection, and server init uses a session bound to it. --- src/dstack/_internal/server/app.py | 17 +++++++++++------ src/dstack/_internal/server/db.py | 25 +++++++++++++++---------- 2 files changed, 26 insertions(+), 16 deletions(-) diff --git a/src/dstack/_internal/server/app.py b/src/dstack/_internal/server/app.py index 0b06dd6b20..904cb3285d 100644 --- a/src/dstack/_internal/server/app.py +++ b/src/dstack/_internal/server/app.py @@ -23,7 +23,7 @@ from dstack._internal.server.background.pipeline_tasks import start_pipeline_tasks from dstack._internal.server.background.scheduled_tasks import start_scheduled_tasks from dstack._internal.server.background.scheduled_tasks.probes import PROBES_SCHEDULER -from dstack._internal.server.db import get_db, get_session_ctx, migrate +from dstack._internal.server.db import get_db, migrate from dstack._internal.server.routers import ( auth, backends, @@ -130,15 +130,19 @@ async def lifespan(app: FastAPI): server_config_loaded = server_config_manager.load_config() # Encryption has to be configured before working with users and projects await server_config_manager.apply_encryption() - async with get_db().engine.connect() as lock_connection: - # The lock is bound to a dedicated connection because the session commits inside the block, - # which would move a session-bound lock release to a different pooled connection. + async with get_db().engine.connect() as connection: + # Running server init using a dedicated connection because there are multiple + # transactions/commits happening under the advisory lock + # and we need to guarantee the same connection releases the lock. async with advisory_lock_ctx( - bind=lock_connection, + bind=connection, dialect_name=get_db().dialect_name, resource="server_init", ): - async with get_session_ctx() as session: + # End the lock statement's transaction so that the session controls its own transactions. + # The session-level lock survives the commit. + await connection.commit() + async with get_db().get_session(bind=connection) as session: admin, _ = await get_or_create_admin_user(session=session) await get_or_create_default_project( session=session, @@ -162,6 +166,7 @@ async def lifespan(app: FastAPI): {"show_path": False}, ) await server_config_manager.apply_config(session=session, owner=admin) + await session.commit() update_default_project( project_name=DEFAULT_PROJECT_NAME, diff --git a/src/dstack/_internal/server/db.py b/src/dstack/_internal/server/db.py index ce9fc82067..de1b21354c 100644 --- a/src/dstack/_internal/server/db.py +++ b/src/dstack/_internal/server/db.py @@ -6,6 +6,7 @@ from sqlalchemy import AsyncAdaptedQueuePool, event, make_url from sqlalchemy.engine.interfaces import DBAPIConnection from sqlalchemy.ext.asyncio import ( + AsyncConnection, AsyncEngine, AsyncSession, async_sessionmaker, @@ -54,8 +55,15 @@ def set_sqlite_pragma(dbapi_connection: DBAPIConnection, _: ConnectionPoolEntry) def dialect_name(self) -> str: return self.engine.dialect.name - def get_session(self) -> AsyncSession: - return self.session_maker() + def get_session(self, bind: Optional[AsyncConnection] = None) -> AsyncSession: + """ + Returns a new session. If `bind` is given, the session runs on that connection + instead of checking connections out of the pool. The connection must not be in + a transaction, otherwise the session joins it and its commits do not commit. + """ + if bind is None: + return self.session_maker() + return self.session_maker(bind=bind) def _get_connect_args(self, url: str) -> dict: if make_url(url).get_backend_name() == "postgresql": @@ -98,20 +106,17 @@ async def migrate(): # The lock is polled instead of awaited in pg_advisory_lock() to avoid this: # migrations waiting for older snapshots to finish (e.g. CREATE INDEX CONCURRENTLY) # wait for the blocked replicas waiting for the "migrations" lock, deadlocking both. - # The lock connection runs in autocommit mode because an idle - # transaction between attempts would keep its snapshot and hang both replicas the same way. - # Migrations run on a separate, transactional connection. - async with db.engine.connect() as lock_connection: - await lock_connection.execution_options(isolation_level="AUTOCOMMIT") + async with db.engine.connect() as connection: for _ in range(_MIGRATIONS_LOCK_MAX_ATTEMPTS): async with try_advisory_lock_ctx( - bind=lock_connection, + bind=connection, dialect_name=db.dialect_name, resource="migrations", ) as locked: + # End the attempt's transaction so that no snapshot is held while waiting. + await connection.commit() if locked: - async with db.engine.connect() as connection: - await connection.run_sync(_run_alembic_upgrade) + await connection.run_sync(_run_alembic_upgrade) return await asyncio.sleep(_MIGRATIONS_LOCK_POLL_INTERVAL) raise TimeoutError( From 06d3513098c89220bba053e8f4cf79d75ac24c7e Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Tue, 15 Sep 2026 14:14:09 +0500 Subject: [PATCH 7/8] Fix advisory lock documentation example --- contributing/LOCKING.md | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/contributing/LOCKING.md b/contributing/LOCKING.md index 604b1a211b..420c493c01 100644 --- a/contributing/LOCKING.md +++ b/contributing/LOCKING.md @@ -119,11 +119,15 @@ Postgres has two kinds of advisory locks: * `pg_advisory_xact_lock` is released when the transaction ends. The unique names pattern above uses it. * `pg_advisory_lock` is bound to the connection. It survives commit and rollback and is released only by `pg_advisory_unlock` on the same connection or when the connection closes. `advisory_lock_ctx()` in `services/locking.py` wraps this kind for work that must span several transactions, such as migrations or server initialization. -Either use `pg_advisory_xact_lock` within a single transaction, or use `advisory_lock_ctx()` and follow these rules: +Within a single transaction, prefer `pg_advisory_xact_lock` for automatic release on commit or rollback. When using `advisory_lock_ctx()`, the requirement is to acquire and release the lock on the same physical connection: -* Bind it to an `AsyncConnection` from `engine.connect()`, not to an `AsyncSession`. If the session commits inside the block, its next statement may run on a different pooled connection, and `pg_advisory_unlock` goes to a connection that never held the lock. Postgres only returns `false` with a warning in this case, so the failure is silent: the lock stays on an idle pooled connection until the process exits, and every replica blocks forever on its next acquire. See https://github.com/dstackai/dstack/issues/3881 for an example. +* An `AsyncSession` is a valid bind when the locked block stays within one transaction. Do not commit, roll back, or close the session before releasing the lock. +* If the block spans multiple transactions, hold an `AsyncConnection` from `engine.connect()` and bind the lock and any sessions to it. Otherwise, a session's next transaction may use a different pooled connection, and `pg_advisory_unlock` goes to a connection that never held the lock. Postgres only returns `false` with a warning in this case, so the failure is silent: the lock stays on an idle pooled connection until the process exits, and every replica blocks forever on its next acquire. See https://github.com/dstackai/dstack/issues/3881 for an example. +* Run protected DB work on the connection holding the lock. When binding a new session to an explicit connection, end the lock statement's transaction first so that the session controls its own transactions. If the connection is lost, abort the protected operation instead of reconnecting and continuing without the lock. * Keep the locked block short and bounded. Every waiter is blocked inside `pg_advisory_lock` holding a DB connection of its own, so a long or hung holder pins one connection per waiter. See `DATABASE.md`. +For a block that spans multiple transactions: + ```python async with get_db().engine.connect() as connection: async with advisory_lock_ctx( @@ -131,9 +135,12 @@ async with get_db().engine.connect() as connection: dialect_name=get_db().dialect_name, resource="server_init", ): - async with get_session_ctx() as session: - # The session may commit freely: the lock lives on `connection`. + # End the lock statement's transaction. The session-level lock survives the commit. + await connection.commit() + async with get_db().get_session(bind=connection) as session: + # Session commits keep using the connection that holds the lock. ... + await session.commit() ``` A released connection goes back to the pool, so a lock that failed to release stays there too. `_release_advisory_lock()` tolerates failures because the common one is an invalidated connection, in which case Postgres has already dropped the lock. A release that fails on a live connection strands the lock. From 9ab5b1605899147505673e52c7334c42531f0362 Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Tue, 15 Sep 2026 14:21:00 +0500 Subject: [PATCH 8/8] Drop extra docstring --- .../server/background/pipeline_tasks/runs/__init__.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/dstack/_internal/server/background/pipeline_tasks/runs/__init__.py b/src/dstack/_internal/server/background/pipeline_tasks/runs/__init__.py index a40b003db4..713828d27b 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/runs/__init__.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/runs/__init__.py @@ -534,10 +534,6 @@ async def _load_active_context( async def _load_gateway_stats(run_model: RunModel, run_spec: RunSpec) -> Optional[PerWindowStats]: - """ - Fetches service stats from the gateway replicas. - Talks to gateways over SSH, so it must run outside DB sessions. - """ if run_spec.configuration.type != "service" or run_model.gateway is None: return None return await get_combined_gateway_stats(