Speed Python up tests with session-scoped db fixtures - #4088
Merged
Conversation
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.
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.
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.
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.
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.
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.
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Significant speedup of the dstack Python test suite. On my MacBook Pro M1:
--runpostgres: ~110s to ~40s.-n autoexecution without--runpostgres: ~70s to ~25s.--runpostgres: >350s to ~95s.-n autoexecution with--runpostgres: ~150s to 50s.The speedup brings all the obvious benefits like faster CI times. My initial motivation was discovering that agents spend >50% of time running the dstack test suite when working on large features/refactoring since they have to run the full test suite. So this also improves agents productivity.
Changes: