From 39027ced9ca886c358a5e5ccbe35626a8c530258 Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Tue, 4 Aug 2026 10:01:57 +0500 Subject: [PATCH 01/11] Fix stale test config paths in AGENTS.md --- AGENTS.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 71e734258..99f4392c3 100644 --- a/AGENTS.md +++ b/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. From 5aa2ef5133d26d547d3ece7127363bdcc548b00b Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Tue, 4 Aug 2026 10:18:29 +0500 Subject: [PATCH 02/11] Speed up tests that wait on real timeouts Three fleet delete tests asserted lock contention errors by letting the retry loop exhaust, costing 5s each, and the gateway timeout test waited out the 8s replica client timeout shared by the whole module. Name the fleet retry bounds so tests can drop the wait, and give the timeout test its own short-timeout fixture; the shared 8s one stays, since httpbin needs it to serve the other tests. Also skip the sleep after the final attempt, which delayed the error response by 500ms in production. --- .../_internal/server/services/fleets.py | 14 ++++++++---- .../_internal/server/routers/test_fleets.py | 13 ++++++++--- .../proxy/routers/test_service_proxy.py | 22 ++++++++++++++++--- 3 files changed, 39 insertions(+), 10 deletions(-) diff --git a/src/dstack/_internal/server/services/fleets.py b/src/dstack/_internal/server/services/fleets.py index 2c6a924c5..d860a4ffd 100644 --- a/src/dstack/_internal/server/services/fleets.py +++ b/src/dstack/_internal/server/services/fleets.py @@ -101,6 +101,10 @@ 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." diff --git a/src/tests/_internal/server/routers/test_fleets.py b/src/tests/_internal/server/routers/test_fleets.py index c8f550779..4cf8ce2b7 100644 --- a/src/tests/_internal/server/routers/test_fleets.py +++ b/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) diff --git a/src/tests/_internal/server/services/proxy/routers/test_service_proxy.py b/src/tests/_internal/server/services/proxy/routers/test_service_proxy.py index cf31e8af0..3df1404b9 100644 --- a/src/tests/_internal/server/services/proxy/routers/test_service_proxy.py +++ b/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" From 589c26f338b81992b7482092d5cfab5525c564d3 Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Tue, 4 Aug 2026 12:16:08 +0500 Subject: [PATCH 03/11] Reuse one SQLite database across tests test_db built a new engine and ran create_all per test, which cost ~20ms each, and a fresh engine also discarded SQLAlchemy's compiled statement cache and re-ran the connect PRAGMAs. Create the schema once per session and clear the rows between tests instead, mirroring what the Postgres path already does. Sharing a session-scoped engine needs a session-scoped event loop, so tests must no longer leave pending tasks behind; none do today. Cuts the serial suite from ~112s to ~57s. --- pyproject.toml | 5 ++ src/dstack/_internal/server/testing/conf.py | 71 +++++++++++++-------- src/tests/conftest.py | 1 + 3 files changed, 51 insertions(+), 26 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index efa7986c0..f6713bd4a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -128,6 +128,11 @@ addopts = [ # unix socket for Docker/testcontainers "--allow-unix-socket", ] +# One event loop for the whole session so session-scoped async fixtures, notably the +# `sqlite_db` engine, stay usable by every test. Tests must not leave pending tasks +# behind: with a shared loop they outlive the test that spawned them. +asyncio_default_fixture_loop_scope = "session" +asyncio_default_test_loop_scope = "session" markers = [ "shim_version", "dockerized", diff --git a/src/dstack/_internal/server/testing/conf.py b/src/dstack/_internal/server/testing/conf.py index 1847b1cd3..2036cd540 100644 --- a/src/dstack/_internal/server/testing/conf.py +++ b/src/dstack/_internal/server/testing/conf.py @@ -18,39 +18,52 @@ def postgres_container(): yield postgres.get_connection_url() -# test_db is function-scoped since making it session-scoped did not bring much benefit. +SQLITE_URL = "sqlite+aiosqlite://" + + +@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 -async def test_db(request): +async def test_db(request, sqlite_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, - ) - elif db_type == "postgres": - if not request.config.getoption("--runpostgres"): - pytest.skip("Skipping Postgres tests as --runpostgres was not provided") - db_url = request.getfixturevalue("postgres_container") - else: + override_db(sqlite_db) + await _delete_all_rows(sqlite_db) + yield sqlite_db + return + if db_type != "postgres": raise ValueError(f"Unknown db_type {db_type}") - db = Database(db_url, engine=engine) + if not request.config.getoption("--runpostgres"): + pytest.skip("Skipping Postgres tests as --runpostgres was not provided") + db_url = request.getfixturevalue("postgres_container") + db = Database(db_url) override_db(db) - if db_type == "sqlite": + if db_url not in _initialized_postgres_db_urls: 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) + _initialized_postgres_db_urls.add(db_url) + await _truncate_postgres_db(db) yield db await db.engine.dispose() @@ -62,6 +75,12 @@ async def session(test_db): yield session +async def _delete_all_rows(db: Database): + async with db.engine.begin() as conn: + for table in reversed(BaseModel.metadata.sorted_tables): + await conn.exec_driver_sql(f'DELETE FROM "{table.name}"') + + async def _truncate_postgres_db(db: Database): preparer = db.engine.sync_engine.dialect.identifier_preparer table_names = ", ".join( diff --git a/src/tests/conftest.py b/src/tests/conftest.py index cef0711de..097bbd1d5 100644 --- a/src/tests/conftest.py +++ b/src/tests/conftest.py @@ -6,6 +6,7 @@ from dstack._internal.server.testing.conf import ( # noqa: F401 postgres_container, session, + sqlite_db, test_db, ) from dstack._internal.settings import FeatureFlags From 11944c50d708fada47dfd929459a9000f04bcd34 Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Tue, 4 Aug 2026 13:18:00 +0500 Subject: [PATCH 04/11] Reuse one Postgres database across tests The Postgres path built a new Database per test and cleaned up with TRUNCATE. TRUNCATE takes an exclusive lock and rewrites files, costing ~90ms for these tables against ~2ms to delete the rows, and a new engine per test also discarded the compiled statement cache. Share the session-wide database and clear rows the same way SQLite does, batching the deletes into one statement since a round trip per table dominates over a socket. Also drop fsync in the container: a test database never has to survive a crash. Two plan tests asserted the order of instances returned by a query that does not order them; TRUNCATE hid this by resetting the heap every test. Cuts the suite with --runpostgres from ~339s to ~147s. --- src/dstack/_internal/server/testing/conf.py | 83 +++++++++++-------- .../server/services/runs/test_plan.py | 6 +- src/tests/conftest.py | 1 + 3 files changed, 53 insertions(+), 37 deletions(-) diff --git a/src/dstack/_internal/server/testing/conf.py b/src/dstack/_internal/server/testing/conf.py index 2036cd540..3ac01f0f4 100644 --- a/src/dstack/_internal/server/testing/conf.py +++ b/src/dstack/_internal/server/testing/conf.py @@ -8,19 +8,20 @@ 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() -SQLITE_URL = "sqlite+aiosqlite://" - - @pytest_asyncio.fixture(scope="session") async def sqlite_db(): """ @@ -44,28 +45,38 @@ async def sqlite_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, sqlite_db): +async def test_db(request, sqlite_db, postgres_db): db_type = getattr(request, "param", "sqlite") if db_type == "sqlite": - override_db(sqlite_db) - await _delete_all_rows(sqlite_db) - yield sqlite_db - return - if db_type != "postgres": + db = sqlite_db + elif db_type == "postgres": + if postgres_db is None: + pytest.skip("Skipping Postgres tests as --runpostgres was not provided") + db = postgres_db + else: raise ValueError(f"Unknown db_type {db_type}") - if not request.config.getoption("--runpostgres"): - pytest.skip("Skipping Postgres tests as --runpostgres was not provided") - db_url = request.getfixturevalue("postgres_container") - db = Database(db_url) override_db(db) - 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 @@ -75,19 +86,21 @@ async def session(test_db): yield session -async def _delete_all_rows(db: Database): - async with db.engine.begin() as conn: - for table in reversed(BaseModel.metadata.sorted_tables): - await conn.exec_driver_sql(f'DELETE FROM "{table.name}"') - +async def _clear_tables(db: Database): + """ + Removes every row, leaving the schema in place. -async def _truncate_postgres_db(db: Database): + 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: - return - truncate_statement = f"TRUNCATE {table_names} RESTART IDENTITY CASCADE" + names = [preparer.format_table(table) for table in reversed(BaseModel.metadata.sorted_tables)] async with db.engine.begin() as conn: - await conn.exec_driver_sql(truncate_statement) + if db.dialect_name == "postgresql": + # Batch into one statement: a round trip per table is most of the cost here. + statements = " ".join(f"DELETE FROM {name};" for name in names) + await conn.exec_driver_sql(f"DO $$ BEGIN {statements} END $$;") + return + for name in names: + await conn.exec_driver_sql(f"DELETE FROM {name}") diff --git a/src/tests/_internal/server/services/runs/test_plan.py b/src/tests/_internal/server/services/runs/test_plan.py index ac8491385..fab3a19e2 100644 --- a/src/tests/_internal/server/services/runs/test_plan.py +++ b/src/tests/_internal/server/services/runs/test_plan.py @@ -503,7 +503,8 @@ 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] + # Compared as a set: the plan query does not order instances. + assert {instance for instance, _ in offers} == {selected_1, selected_2} assert [offer.blocks for _, offer in offers] == [2, 2] assert [offer.total_blocks for _, offer in offers] == [2, 2] @@ -555,7 +556,8 @@ 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] + # Compared as a set: the plan query does not order instances. + assert {instance for instance, _ in offers} == {selected_1, selected_2} @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) diff --git a/src/tests/conftest.py b/src/tests/conftest.py index 097bbd1d5..5d7481d24 100644 --- a/src/tests/conftest.py +++ b/src/tests/conftest.py @@ -5,6 +5,7 @@ from dstack._internal.server.testing.conf import ( # noqa: F401 postgres_container, + postgres_db, session, sqlite_db, test_db, From 8a1790946cc15dc6009f5770675c13629f4d22bb Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Tue, 4 Aug 2026 14:18:39 +0500 Subject: [PATCH 05/11] Assert plan offers by sorted ids, not as a set --- .../_internal/server/services/runs/test_plan.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/src/tests/_internal/server/services/runs/test_plan.py b/src/tests/_internal/server/services/runs/test_plan.py index fab3a19e2..df457d33f 100644 --- a/src/tests/_internal/server/services/runs/test_plan.py +++ b/src/tests/_internal/server/services/runs/test_plan.py @@ -503,8 +503,10 @@ async def test_multinode_returns_full_host_offer_per_selected_shared_instance( exclude_not_available=True, ) - # Compared as a set: the plan query does not order instances. - 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] @@ -556,8 +558,10 @@ async def test_multinode_returns_selected_instances_in_same_cluster_fleet( exclude_not_available=True, ) - # Compared as a set: the plan query does not order instances. - 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) From 11d2e2798b34e790d45bcc97094b37a37e7b86ac Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Tue, 4 Aug 2026 14:28:27 +0500 Subject: [PATCH 06/11] Explain the session event loop scope in pytest config --- pyproject.toml | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index f6713bd4a..b2ce011a8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -128,9 +128,8 @@ addopts = [ # unix socket for Docker/testcontainers "--allow-unix-socket", ] -# One event loop for the whole session so session-scoped async fixtures, notably the -# `sqlite_db` engine, stay usable by every test. Tests must not leave pending tasks -# behind: with a shared loop they outlive the test that spawned them. +# 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 = [ From e7f9dd303aab100d0ffc204a7120d7653129d7bc Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Tue, 4 Aug 2026 15:26:47 +0500 Subject: [PATCH 07/11] Generate one RSA key pair per test session Tests generated 104 key pairs costing 7.7s: the server generates one per user, project, gateway, and job, and a 2048-bit key takes ~80ms. Reuse a single pair for the whole session. Route the two remaining call sites through the `crypto` module so all four resolve the function at call time and one patch covers them. --- .../server/services/gateways/__init__.py | 4 ++-- .../_internal/server/services/projects.py | 4 ++-- src/tests/conftest.py | 20 +++++++++++++++++++ 3 files changed, 24 insertions(+), 4 deletions(-) diff --git a/src/dstack/_internal/server/services/gateways/__init__.py b/src/dstack/_internal/server/services/gateways/__init__.py index af29e21fd..d2fbadc5c 100644 --- a/src/dstack/_internal/server/services/gateways/__init__.py +++ b/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() diff --git a/src/dstack/_internal/server/services/projects.py b/src/dstack/_internal/server/services/projects.py index 8d7ad534b..31f995ab8 100644 --- a/src/dstack/_internal/server/services/projects.py +++ b/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(), diff --git a/src/tests/conftest.py b/src/tests/conftest.py index 5d7481d24..d24d1be08 100644 --- a/src/tests/conftest.py +++ b/src/tests/conftest.py @@ -11,6 +11,7 @@ test_db, ) from dstack._internal.settings import FeatureFlags +from dstack._internal.utils import crypto def pytest_configure(config): @@ -64,6 +65,25 @@ def pytest_collection_modifyitems(config, items): item.add_marker(skip_posix) +@pytest.fixture(scope="session", autouse=True) +def reuse_one_rsa_key_pair(): + """ + Hands the same RSA key pair to every caller for the whole test session. + + Generating a 2048-bit key takes ~80ms and the server generates one per user, project, + gateway, and job. A test that needs two different keys has to generate its own. + """ + private_bytes, public_bytes = crypto.generate_rsa_key_pair_bytes() + public_key = public_bytes.rsplit(b" ", 1)[0] # drop the comment, callers pass their own + + def generate_rsa_key_pair_bytes(comment: str = "dstack") -> tuple[bytes, bytes]: + return private_bytes, public_key + f" {comment}\n".encode() + + with pytest.MonkeyPatch.context() as monkeypatch: + monkeypatch.setattr(crypto, "generate_rsa_key_pair_bytes", generate_rsa_key_pair_bytes) + yield + + @pytest.fixture(scope="session", autouse=True) def disable_feature_flags(): """ From 513602c933e656cf84b0ec2e8d358764a28d25d2 Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Tue, 4 Aug 2026 15:26:58 +0500 Subject: [PATCH 08/11] Batch the per-test row deletes into one round trip Clearing 32 tables with a statement each cost ~3.4ms per test against ~0.8ms sent together. SQLite has no multi-statement execute, so use `executescript` on the driver connection; Postgres already batched via a DO block. --- src/dstack/_internal/server/testing/conf.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/src/dstack/_internal/server/testing/conf.py b/src/dstack/_internal/server/testing/conf.py index 3ac01f0f4..e1dabfed2 100644 --- a/src/dstack/_internal/server/testing/conf.py +++ b/src/dstack/_internal/server/testing/conf.py @@ -93,14 +93,19 @@ async def _clear_tables(db: Database): 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. + + Both dialects send the deletes in one round trip, which is most of the remaining cost: + 32 separate statements take ~3.4ms against ~0.8ms batched. """ preparer = db.engine.sync_engine.dialect.identifier_preparer names = [preparer.format_table(table) for table in reversed(BaseModel.metadata.sorted_tables)] - async with db.engine.begin() as conn: - if db.dialect_name == "postgresql": - # Batch into one statement: a round trip per table is most of the cost here. - statements = " ".join(f"DELETE FROM {name};" for name in names) + 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 - for name in names: - await conn.exec_driver_sql(f"DELETE FROM {name}") + return + 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) From 44356537b3b423926ba190062aebd3d25e7cd3b6 Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Tue, 4 Aug 2026 15:28:03 +0500 Subject: [PATCH 09/11] Update the docstring --- src/dstack/_internal/server/testing/conf.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/dstack/_internal/server/testing/conf.py b/src/dstack/_internal/server/testing/conf.py index e1dabfed2..19aa51d53 100644 --- a/src/dstack/_internal/server/testing/conf.py +++ b/src/dstack/_internal/server/testing/conf.py @@ -93,9 +93,6 @@ async def _clear_tables(db: Database): 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. - - Both dialects send the deletes in one round trip, which is most of the remaining cost: - 32 separate statements take ~3.4ms against ~0.8ms batched. """ preparer = db.engine.sync_engine.dialect.identifier_preparer names = [preparer.format_table(table) for table in reversed(BaseModel.metadata.sorted_tables)] From 6f9b2d489dd032cd9d38f73472bbbd906248b65b Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Tue, 4 Aug 2026 15:32:49 +0500 Subject: [PATCH 10/11] Parse each SSH key once per test session Fleet spec validation parsed the same test keys 20 times at ~190ms each, 3.8s in total, because paramiko validates an RSA key on load. Cache the parses for the session, and route the call site through the `ssh` module so one patch covers it. --- src/dstack/_internal/server/services/fleets.py | 4 ++-- src/tests/_internal/server/conftest.py | 15 +++++++++++++++ 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/src/dstack/_internal/server/services/fleets.py b/src/dstack/_internal/server/services/fleets.py index d860a4ffd..0173b944d 100644 --- a/src/dstack/_internal/server/services/fleets.py +++ b/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,7 +98,6 @@ 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__) @@ -1440,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. " diff --git a/src/tests/_internal/server/conftest.py b/src/tests/_internal/server/conftest.py index 125cc5de1..b46e2fdcc 100644 --- a/src/tests/_internal/server/conftest.py +++ b/src/tests/_internal/server/conftest.py @@ -1,4 +1,5 @@ from collections.abc import Generator +from functools import cache from pathlib import Path from unittest.mock import AsyncMock, Mock, patch @@ -15,6 +16,7 @@ session, test_db, ) +from dstack._internal.utils import ssh as ssh_utils def _warm_up_route_schemas() -> None: @@ -44,6 +46,19 @@ def _warm_up_route_schemas() -> None: _warm_up_route_schemas() +@pytest.fixture(scope="session", autouse=True) +def cache_parsed_ssh_keys(): + """ + Parses each SSH key once per session instead of once per call. + + Parsing an RSA key costs ~190ms because paramiko validates it, and fleet spec + validation parses the same handful of test keys over and over. + """ + with pytest.MonkeyPatch.context() as monkeypatch: + monkeypatch.setattr(ssh_utils, "pkey_from_str", cache(ssh_utils.pkey_from_str)) + yield + + @pytest.fixture def client(): transport = httpx.ASGITransport(app=app) From 938be8312d8c00255e43f41cb11e264ae9f2a912 Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Tue, 4 Aug 2026 15:45:42 +0500 Subject: [PATCH 11/11] Keep both test key fixtures in one conftest The SSH key fixture sat in the server conftest to avoid importing paramiko for non-server runs, but the root conftest already pulls it in transitively, so there was nothing to avoid. Placing both together also covers the CLI fleet configurator and drops the order-dependent patching. --- src/tests/_internal/server/conftest.py | 15 --------------- src/tests/conftest.py | 15 +++++++++++++++ 2 files changed, 15 insertions(+), 15 deletions(-) diff --git a/src/tests/_internal/server/conftest.py b/src/tests/_internal/server/conftest.py index b46e2fdcc..125cc5de1 100644 --- a/src/tests/_internal/server/conftest.py +++ b/src/tests/_internal/server/conftest.py @@ -1,5 +1,4 @@ from collections.abc import Generator -from functools import cache from pathlib import Path from unittest.mock import AsyncMock, Mock, patch @@ -16,7 +15,6 @@ session, test_db, ) -from dstack._internal.utils import ssh as ssh_utils def _warm_up_route_schemas() -> None: @@ -46,19 +44,6 @@ def _warm_up_route_schemas() -> None: _warm_up_route_schemas() -@pytest.fixture(scope="session", autouse=True) -def cache_parsed_ssh_keys(): - """ - Parses each SSH key once per session instead of once per call. - - Parsing an RSA key costs ~190ms because paramiko validates it, and fleet spec - validation parses the same handful of test keys over and over. - """ - with pytest.MonkeyPatch.context() as monkeypatch: - monkeypatch.setattr(ssh_utils, "pkey_from_str", cache(ssh_utils.pkey_from_str)) - yield - - @pytest.fixture def client(): transport = httpx.ASGITransport(app=app) diff --git a/src/tests/conftest.py b/src/tests/conftest.py index d24d1be08..4bedf948b 100644 --- a/src/tests/conftest.py +++ b/src/tests/conftest.py @@ -1,5 +1,6 @@ import inspect import os +from functools import cache import pytest @@ -12,6 +13,7 @@ ) from dstack._internal.settings import FeatureFlags from dstack._internal.utils import crypto +from dstack._internal.utils import ssh as ssh_utils def pytest_configure(config): @@ -84,6 +86,19 @@ def generate_rsa_key_pair_bytes(comment: str = "dstack") -> tuple[bytes, bytes]: yield +@pytest.fixture(scope="session", autouse=True) +def cache_parsed_ssh_keys(): + """ + Parses each SSH key once per session instead of once per call. + + Parsing an RSA key costs ~190ms because paramiko validates it, and fleet spec + validation parses the same handful of test keys over and over. + """ + with pytest.MonkeyPatch.context() as monkeypatch: + monkeypatch.setattr(ssh_utils, "pkey_from_str", cache(ssh_utils.pkey_from_str)) + yield + + @pytest.fixture(scope="session", autouse=True) def disable_feature_flags(): """