Speed Python up tests with session-scoped db fixtures

AGENTS.md

@@ -34,11 +34,11 @@ Before touching a subsystem, read the relevant notes in `contributing/`: `ARCHIT
 - Never edit a migration that has already been applied or released; add a new migration instead.
 
 ## Testing Guidelines
-- Default to `uv run pytest`. Use markers from `tests/conftest.py` like `--runpostgres` if need to include specific tests.
+- Default to `uv run pytest`. Use markers from `src/tests/conftest.py` like `--runpostgres` if need to include specific tests.
 - Scope the run to the change: for trivial or localized edits, run only the affected test modules, `Test*` classes, or `-k` selection instead of the whole suite. Reserve the full suite for broad or cross-cutting changes.
 - Speed up large runs with `-n auto` (pytest-xdist), e.g. `uv run pytest -n auto`.
 - Group tests for the same unit (function/class) using `Test*` classes that mirror unit's name.
-- Keep tests hermetic (network disabled except localhost per `pytest.ini`); stub cloud calls with mocks.
+- Keep tests hermetic (network disabled except localhost per `[tool.pytest.ini_options]` in `pyproject.toml`); stub cloud calls with mocks.
 
 ## Commit & Pull Request Guidelines
 - Name branches `issue_{issue_num}_{title}` when the work tracks an issue (e.g. `issue_3959_replicated_alb_gateways`), and `pr_{title}` otherwise.

pyproject.toml

@@ -128,6 +128,10 @@ addopts = [
     # unix socket for Docker/testcontainers
     "--allow-unix-socket",
 ]
+# One session-scoped event loop for fixtures and tests so that `sqlite_db` and `postgres_db` engines are reused.
+# Tests must not leave pending asyncio tasks behind; nothing enforces it automatically.
+asyncio_default_fixture_loop_scope = "session"
+asyncio_default_test_loop_scope = "session"
 markers = [
     "shim_version",
     "dockerized",

src/dstack/_internal/server/services/fleets.py

@@ -89,6 +89,7 @@
 )
 from dstack._internal.server.services.resources import set_resources_defaults
 from dstack._internal.utils import random_names
+from dstack._internal.utils import ssh as ssh_utils
 from dstack._internal.utils.common import (
     EntityID,
     EntityName,
@@ -97,10 +98,13 @@
     get_lowest_unused_nums,
 )
 from dstack._internal.utils.logging import get_logger
-from dstack._internal.utils.ssh import pkey_from_str
 
 logger = get_logger(__name__)
 
+# How hard to retry row locks before reporting the fleet as busy.
+_LOCK_RETRY_ATTEMPTS = 10
+_LOCK_RETRY_INTERVAL = 0.5
+
 
 def switch_fleet_status(
     session: AsyncSession,
@@ -788,7 +792,7 @@ async def delete_fleets(
         # Retry locking fleets to increase lock acquisition chances.
         # This hack is needed until requests are queued.
         fleet_models = []
-        for i in range(10):
+        for attempt in range(_LOCK_RETRY_ATTEMPTS):
             res = await session.execute(
                 select(FleetModel)
                 .where(
@@ -814,7 +818,8 @@ async def delete_fleets(
             fleet_models = res.scalars().unique().all()
             if len(fleet_models) == len(fleets_ids):
                 break
-            await asyncio.sleep(0.5)
+            if attempt < _LOCK_RETRY_ATTEMPTS - 1:
+                await asyncio.sleep(_LOCK_RETRY_INTERVAL)
         if len(fleet_models) != len(fleets_ids):
             # TODO: Make the endpoint fully async so we don't need to lock and error.
             msg = (
@@ -826,7 +831,7 @@ async def delete_fleets(
         # Retry locking instances to increase lock acquisition chances.
         # This hack is needed until requests are queued.
         instances_left_to_lock = set(instances_ids)
-        for i in range(10):
+        for attempt in range(_LOCK_RETRY_ATTEMPTS):
             res = await session.execute(
                 select(InstanceModel.id)
                 .where(
@@ -841,7 +846,8 @@ async def delete_fleets(
             instances_left_to_lock.difference_update(res.scalars().unique().all())
             if len(instances_left_to_lock) == 0:
                 break
-            await asyncio.sleep(0.5)
+            if attempt < _LOCK_RETRY_ATTEMPTS - 1:
+                await asyncio.sleep(_LOCK_RETRY_INTERVAL)
         if len(instances_left_to_lock) > 0:
             msg = (
                 "Failed to delete fleets: fleet instances are being processed currently. Try again later."
@@ -1434,7 +1440,7 @@ def _validate_ssh_key(ssh_key: SSHKey):
     if ssh_key.private is None:
         raise ServerClientError("Private key not provided")
     try:
-        pkey_from_str(ssh_key.private)
+        ssh_utils.pkey_from_str(ssh_key.private)
     except ValueError:
         raise ServerClientError(
             "Unsupported key type. "

src/dstack/_internal/server/services/gateways/__init__.py

@@ -84,12 +84,12 @@
 from dstack._internal.server.services.plugins import apply_plugin_policies
 from dstack._internal.server.utils.common import gather_map_async
 from dstack._internal.settings import FeatureFlags
+from dstack._internal.utils import crypto
 from dstack._internal.utils.common import (
     get_current_datetime,
     get_or_error,
     interpolate_gateway_domain,
 )
-from dstack._internal.utils.crypto import generate_rsa_key_pair_bytes
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
@@ -198,7 +198,7 @@ def create_gateway_compute_model(
 ) -> GatewayComputeModel:
     assert configuration.name is not None
 
-    private_bytes, public_bytes = generate_rsa_key_pair_bytes()
+    private_bytes, public_bytes = crypto.generate_rsa_key_pair_bytes()
     gateway_ssh_private_key = private_bytes.decode()
     gateway_ssh_public_key = public_bytes.decode()
 

src/dstack/_internal/server/services/projects.py

@@ -50,8 +50,8 @@
 )
 from dstack._internal.server.services.permissions import get_default_permissions
 from dstack._internal.server.settings import DEFAULT_PROJECT_NAME
+from dstack._internal.utils import crypto
 from dstack._internal.utils.common import get_current_datetime, run_async
-from dstack._internal.utils.crypto import generate_rsa_key_pair_bytes
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
@@ -628,7 +628,7 @@ async def create_project_model(
     validate_project_name(project_name)
     templates_repo = await _normalize_templates_repo_url(templates_repo)
     private_bytes, public_bytes = await run_async(
-        generate_rsa_key_pair_bytes, f"{project_name}@dstack"
+        crypto.generate_rsa_key_pair_bytes, f"{project_name}@dstack"
     )
     project = ProjectModel(
         id=uuid.uuid4(),

src/dstack/_internal/server/testing/conf.py

@@ -8,51 +8,75 @@
 from dstack._internal.server.db import Database, override_db
 from dstack._internal.server.models import BaseModel
 
-# Remember initialized URLs to create metadata once per session.
-_initialized_postgres_db_urls = set()
+SQLITE_URL = "sqlite+aiosqlite://"
 
 
 @pytest.fixture(scope="session")
 def postgres_container():
-    with PostgresContainer("postgres:16-alpine", driver="asyncpg") as postgres:
+    with PostgresContainer(
+        "postgres:16-alpine",
+        driver="asyncpg",
+        # A test database never has to survive a crash, and fsync dominates commit cost.
+        command="postgres -c fsync=off -c synchronous_commit=off -c full_page_writes=off",
+    ) as postgres:
         yield postgres.get_connection_url()
 
 
-# test_db is function-scoped since making it session-scoped did not bring much benefit.
+@pytest_asyncio.fixture(scope="session")
+async def sqlite_db():
+    """
+    A SQLite database with the schema created once for the whole session.
+
+    Creating the schema costs ~20ms, so `test_db` clears the rows between tests
+    instead of rebuilding the database per test.
+    """
+    engine = create_async_engine(
+        SQLITE_URL,
+        echo=settings.SQL_ECHO_ENABLED,
+        # For SQLite, allow accessing the in-memory DB from multiple threads:
+        # https://docs.sqlalchemy.org/en/13/dialects/sqlite.html#using-a-memory-database-in-multiple-threads
+        connect_args={"check_same_thread": False},
+        poolclass=StaticPool,
+    )
+    db = Database(SQLITE_URL, engine=engine)
+    async with engine.begin() as conn:
+        await conn.run_sync(BaseModel.metadata.create_all)
+    yield db
+    await engine.dispose()
+
+
+@pytest_asyncio.fixture(scope="session")
+async def postgres_db(request):
+    """
+    A Postgres database with the schema created once for the whole session.
+
+    Yields `None` without `--runpostgres` so that the container only starts when
+    Postgres tests actually run.
+    """
+    if not request.config.getoption("--runpostgres"):
+        yield None
+        return
+    db = Database(request.getfixturevalue("postgres_container"))
+    async with db.engine.begin() as conn:
+        await conn.run_sync(BaseModel.metadata.create_all)
+    yield db
+    await db.engine.dispose()
+
+
 @pytest_asyncio.fixture
-async def test_db(request):
+async def test_db(request, sqlite_db, postgres_db):
     db_type = getattr(request, "param", "sqlite")
-    engine = None
     if db_type == "sqlite":
-        db_url = "sqlite+aiosqlite://"
-        # For SQLite, allow accessing the in-memory DB from multiple threads:
-        # https://docs.sqlalchemy.org/en/13/dialects/sqlite.html#using-a-memory-database-in-multiple-threads
-        engine = create_async_engine(
-            db_url,
-            echo=settings.SQL_ECHO_ENABLED,
-            connect_args={"check_same_thread": False},
-            poolclass=StaticPool,
-        )
+        db = sqlite_db
     elif db_type == "postgres":
-        if not request.config.getoption("--runpostgres"):
+        if postgres_db is None:
             pytest.skip("Skipping Postgres tests as --runpostgres was not provided")
-        db_url = request.getfixturevalue("postgres_container")
+        db = postgres_db
     else:
         raise ValueError(f"Unknown db_type {db_type}")
-    db = Database(db_url, engine=engine)
     override_db(db)
-    if db_type == "sqlite":
-        async with db.engine.begin() as conn:
-            await conn.run_sync(BaseModel.metadata.create_all)
-            # Relying on function-scoped engine for a clean DB
-    else:
-        if db_url not in _initialized_postgres_db_urls:
-            async with db.engine.begin() as conn:
-                await conn.run_sync(BaseModel.metadata.create_all)
-            _initialized_postgres_db_urls.add(db_url)
-        await _truncate_postgres_db(db)
+    await _clear_tables(db)
     yield db
-    await db.engine.dispose()
 
 
 @pytest_asyncio.fixture
@@ -62,13 +86,23 @@ async def session(test_db):
         yield session
 
 
-async def _truncate_postgres_db(db: Database):
+async def _clear_tables(db: Database):
+    """
+    Removes every row, leaving the schema in place.
+
+    Deletes in reverse dependency order so foreign keys stay satisfied. `DELETE` rather
+    than `TRUNCATE` because Postgres takes an exclusive lock and rewrites files per
+    `TRUNCATE`, which costs ~90ms for these tables against ~2ms to delete the rows.
+    """
     preparer = db.engine.sync_engine.dialect.identifier_preparer
-    table_names = ", ".join(
-        preparer.format_table(table) for table in BaseModel.metadata.sorted_tables
-    )
-    if not table_names:
+    names = [preparer.format_table(table) for table in reversed(BaseModel.metadata.sorted_tables)]
+    statements = "".join(f"DELETE FROM {name};" for name in names)
+    if db.dialect_name == "postgresql":
+        async with db.engine.begin() as conn:
+            await conn.exec_driver_sql(f"DO $$ BEGIN {statements} END $$;")
         return
-    truncate_statement = f"TRUNCATE {table_names} RESTART IDENTITY CASCADE"
-    async with db.engine.begin() as conn:
-        await conn.exec_driver_sql(truncate_statement)
+    async with db.engine.connect() as conn:
+        # SQLite has no multi-statement execute, and `executescript` commits on its own.
+        aiosqlite_connection = (await conn.get_raw_connection()).driver_connection
+        assert aiosqlite_connection is not None
+        await aiosqlite_connection.executescript(statements)

src/tests/_internal/server/routers/test_fleets.py

@@ -32,6 +32,7 @@
 from dstack._internal.core.models.profiles import Profile
 from dstack._internal.core.models.users import GlobalRole, ProjectRole
 from dstack._internal.server.models import FleetModel, InstanceModel
+from dstack._internal.server.services import fleets as fleets_services
 from dstack._internal.server.services.fleets import fleet_model_to_fleet
 from dstack._internal.server.services.permissions import DefaultPermissions
 from dstack._internal.server.services.projects import add_project_member
@@ -1683,6 +1684,12 @@ async def test_importer_member_cannot_apply_plan_on_imported_fleet(
         assert response.status_code == 403
 
 
+@pytest.fixture
+def no_lock_retry_wait(monkeypatch: pytest.MonkeyPatch):
+    """Makes tests asserting lock contention errors exhaust the retries without waiting."""
+    monkeypatch.setattr(fleets_services, "_LOCK_RETRY_INTERVAL", 0)
+
+
 class TestDeleteFleets:
     @pytest.mark.asyncio
     async def test_returns_40x_if_not_authenticated(self, client: AsyncClient):
@@ -1794,7 +1801,7 @@ async def test_returns_400_when_fleet_instance_in_use(
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_returns_400_when_fleet_locked(
-        self, test_db, session: AsyncSession, client: AsyncClient
+        self, test_db, session: AsyncSession, client: AsyncClient, no_lock_retry_wait
     ):
         user = await create_user(session, global_role=GlobalRole.USER)
         project = await create_project(session)
@@ -2001,7 +2008,7 @@ async def test_ignores_lock_on_non_selected_instances(
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_returns_400_when_selected_instance_locked(
-        self, test_db, session: AsyncSession, client: AsyncClient
+        self, test_db, session: AsyncSession, client: AsyncClient, no_lock_retry_wait
     ):
         user = await create_user(session, global_role=GlobalRole.USER)
         project = await create_project(session)
@@ -2086,7 +2093,7 @@ async def test_returns_400_when_deleting_busy_instances(
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_returns_400_when_fleet_locked(
-        self, test_db, session: AsyncSession, client: AsyncClient
+        self, test_db, session: AsyncSession, client: AsyncClient, no_lock_retry_wait
     ):
         user = await create_user(session, global_role=GlobalRole.USER)
         project = await create_project(session)

src/tests/_internal/server/services/proxy/routers/test_service_proxy.py

@@ -19,6 +19,9 @@
 from dstack._internal.server.services.proxy.routers.service_proxy import router
 
 MOCK_REPLICA_CLIENT_TIMEOUT = 8
+# Kept well below `MOCK_REPLICA_CLIENT_TIMEOUT` so the gateway timeout test does not wait 8s.
+# Only the timeout test may use it: httpbin needs more than this to serve the other tests.
+SHORT_REPLICA_CLIENT_TIMEOUT = 0.5
 
 # Using GatewayProxyRepo for tests because it is easier to populate than ServerProxyRepo
 ProxyTestRepo = GatewayProxyRepo
@@ -37,6 +40,19 @@ def mock_replica_client_httpbin(httpbin) -> Generator[None, None, None]:
         yield
 
 
+@pytest.fixture
+def mock_replica_client_httpbin_short_timeout(httpbin) -> Generator[None, None, None]:
+    """Same as `mock_replica_client_httpbin`, but times out quickly"""
+
+    with patch(
+        "dstack._internal.proxy.lib.services.service_connection.ServiceConnectionPool.get_or_add"
+    ) as add_connection_mock:
+        add_connection_mock.return_value.client.return_value = ServiceClient(
+            base_url=httpbin.url, timeout=SHORT_REPLICA_CLIENT_TIMEOUT
+        )
+        yield
+
+
 @pytest.fixture
 def mock_replica_client_path_reporter() -> Generator[None, None, None]:
     """Mocks deployed services. Replaces them with an app that returns the requested path"""
@@ -155,13 +171,13 @@ async def test_proxy_not_leaks_cookies(mock_replica_client_httpbin) -> None:
 
 
 @pytest.mark.asyncio
-async def test_proxy_gateway_timeout(mock_replica_client_httpbin) -> None:
+async def test_proxy_gateway_timeout(mock_replica_client_httpbin_short_timeout) -> None:
     repo = ProxyTestRepo()
     await repo.set_project(make_project("test-proj"))
     await repo.set_service(make_service("test-proj", "httpbin"))
     _, client = make_app_client(repo)
-    assert MOCK_REPLICA_CLIENT_TIMEOUT < 10
-    resp = await client.get("http://test-host/proxy/services/test-proj/httpbin/delay/10")
+    assert SHORT_REPLICA_CLIENT_TIMEOUT < 2
+    resp = await client.get("http://test-host/proxy/services/test-proj/httpbin/delay/2")
     assert resp.status_code == 504
     assert resp.json()["detail"] == "Timed out requesting upstream"
 

src/tests/_internal/server/services/runs/test_plan.py

@@ -503,7 +503,10 @@ async def test_multinode_returns_full_host_offer_per_selected_shared_instance(
             exclude_not_available=True,
         )
 
-        assert [instance for instance, _ in offers] == [selected_1, selected_2]
+        # Sorted: the plan query does not order instances.
+        assert sorted(instance.id for instance, _ in offers) == sorted(
+            [selected_1.id, selected_2.id]
+        )
         assert [offer.blocks for _, offer in offers] == [2, 2]
         assert [offer.total_blocks for _, offer in offers] == [2, 2]
 
@@ -555,7 +558,10 @@ async def test_multinode_returns_selected_instances_in_same_cluster_fleet(
             exclude_not_available=True,
         )
 
-        assert [instance for instance, _ in offers] == [selected_1, selected_2]
+        # Sorted: the plan query does not order instances.
+        assert sorted(instance.id for instance, _ in offers) == sorted(
+            [selected_1.id, selected_2.id]
+        )
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)