Skip to content

fix(tests): de-flake serve tests by isolating ports and failing fast - #2298

Draft
bhimrazy wants to merge 2 commits into
Lightning-AI:mainfrom
bhimrazy:fix/flaky-serve-tests
Draft

fix(tests): de-flake serve tests by isolating ports and failing fast#2298
bhimrazy wants to merge 2 commits into
Lightning-AI:mainfrom
bhimrazy:fix/flaky-serve-tests

Conversation

@bhimrazy

@bhimrazy bhimrazy commented Aug 14, 2026

Copy link
Copy Markdown
Collaborator

What does this PR do?

Fixes the recurring flake in tests/test_serve.py, most often seen on test_serve_with_generate_strategy[tensor_parallel]:

AssertionError: Server did not respond as expected. Error: HTTPConnectionPool(host='127.0.0.1', port=8000):
Max retries exceeded ... Connection refused

Test-only change; nothing under litgpt/ is touched.

Root cause

All seven serve tests bind the hardcoded port 8000, and teardown never waits for the server to actually exit — kill_process_tree only sends the signals, and the thread wrapping the non-blocking Popen was joined instead of the process. That leads to three compounding problems:

  1. Port still held. A server that hasn't finished shutting down makes the next one exit at bind_socket() via uvicorn's silent sys.exit(1), which the test can only observe as a connection refused.
  2. Leaked servers. When an assertion fails, the kill_process_tree call below it is never reached, so the server survives for the rest of the session and takes port 8000 with it. tensor_parallel runs last, so it inherits the debris.
  3. No diagnostics. Nothing checked whether the server was still alive, so a crash at second one looked exactly like a slow start, and the failure message said nothing about why.

Changes

  • Serve each test on its own OS-assigned free port.
  • Poll process.poll() while waiting, and fail immediately with the exit code plus a tail of the server output.
  • Always tear down via a context manager, and wait for the full process tree to exit before the next test starts.
  • Write server output to a file instead of an unread pipe, which can fill up and block a chatty server (NCCL_DEBUG=INFO on CI) before it binds.
  • Fold the boilerplate duplicated across all seven tests into two helpers.

Verification

  • ruff check and ruff format --check clean; pytest tests/test_serve.py passes locally (GPU-gated tests skip).
  • Fail-fast path confirmed by hand: an invalid CLI flag now reports Server exited with code 2 before it was ready. with the server log, instead of a bare connection refused after 30s.

Left as a draft: the 2-GPU tensor_parallel path can't be reproduced locally, so a few CI runs are the real confirmation.

The serve tests all bound the hardcoded port 8000 and killed the server without
waiting for the process tree to exit, so a server that was still shutting down
made the next one exit at bind time with 'Address already in use'. The test then
polled a dead server for 30s and reported a bare connection error with no log.
A failing test leaked its server entirely, because the assertion fired before
the line that killed it.

- serve on a free port per test
- stop polling as soon as the server process exits, and report its exit code
  and log tail instead of a generic connection error
- always kill the server, and wait for its process tree to exit
- log to a file rather than to a pipe that nothing reads during startup
- allow 120s for startup, which is cheap now that a dead server fails fast
- drop the thread wrapper around the non-blocking Popen and share the
  checkpoint and server setup between the tests
Keep reporting the actual exception text (e.g. connection refused) when the
server never answers, instead of a generic status placeholder, matching what
the old _wait_and_check_response reported.
@bhimrazy bhimrazy changed the title fix(tests): stop flaky serve tests from sharing port 8000 fix(tests): de-flake serve tests by isolating ports and failing fast Aug 14, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant