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. diff --git a/pyproject.toml b/pyproject.toml index efa7986c0..b2ce011a8 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/dstack/_internal/server/services/fleets.py b/src/dstack/_internal/server/services/fleets.py index 2c6a924c5..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,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. " 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/dstack/_internal/server/testing/conf.py b/src/dstack/_internal/server/testing/conf.py index 1847b1cd3..19aa51d53 100644 --- a/src/dstack/_internal/server/testing/conf.py +++ b/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) 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" diff --git a/src/tests/_internal/server/services/runs/test_plan.py b/src/tests/_internal/server/services/runs/test_plan.py index ac8491385..df457d33f 100644 --- a/src/tests/_internal/server/services/runs/test_plan.py +++ b/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) diff --git a/src/tests/conftest.py b/src/tests/conftest.py index cef0711de..4bedf948b 100644 --- a/src/tests/conftest.py +++ b/src/tests/conftest.py @@ -1,14 +1,19 @@ import inspect import os +from functools import cache import pytest from dstack._internal.server.testing.conf import ( # noqa: F401 postgres_container, + postgres_db, session, + sqlite_db, test_db, ) 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): @@ -62,6 +67,38 @@ 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 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(): """