[TRTLLM-13409][fix] make proxy shutdown non-blocking when the engine is dead#16312
[TRTLLM-13409][fix] make proxy shutdown non-blocking when the engine is dead#16312JunyiXu-nv wants to merge 5 commits into
Conversation
📝 WalkthroughWalkthroughThe proxy now receives worker PID metadata, monitors local worker liveness with ChangesWorker liveness monitoring
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unittest/executor/test_proxy_fast_death.py (1)
392-456: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCoverage verdict: sufficient for the liveness/registration surface; one gap for follow-up.
These tests cover
_worker_process_alive(kill+reap, recycled-PID),_check_worker_pid_liveness(fast-death path + no-op on empty/doing_shutdown), and_register_worker_process_infos(local-only filter,Noneno-op) well.Gap: the dead-engine bounded-shutdown behavior added in
tensorrt_llm/executor/proxy.py(the_engine_dead-gatedconcurrent.futures.wait(..., timeout=5.0), boundeddispatch_result_thread.join, andmpi_session.shutdown(wait=not self._engine_dead)) has no unit coverage. Consider a test with a fake session/futures asserting thatshutdown()returns within the bound and does not block on pending futures when_engine_deadis set. Reasonable as a follow-up if out of scope here.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/executor/test_proxy_fast_death.py` around lines 392 - 456, Extend the proxy unit tests with coverage for the dead-engine shutdown path in the relevant shutdown method of the proxy. Use fake pending futures and a fake MPI session to verify that when _engine_dead is true, concurrent.futures.wait uses the bounded timeout, dispatch_result_thread.join is bounded, and mpi_session.shutdown receives wait=False without blocking on unfinished work.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/unittest/executor/test_proxy_fast_death.py`:
- Around line 392-456: Extend the proxy unit tests with coverage for the
dead-engine shutdown path in the relevant shutdown method of the proxy. Use fake
pending futures and a fake MPI session to verify that when _engine_dead is true,
concurrent.futures.wait uses the bounded timeout, dispatch_result_thread.join is
bounded, and mpi_session.shutdown receives wait=False without blocking on
unfinished work.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f0bc436d-3971-4813-b39f-9fb292b26582
📒 Files selected for processing (3)
tensorrt_llm/executor/proxy.pytensorrt_llm/executor/worker.pytests/unittest/executor/test_proxy_fast_death.py
|
/bot run --disable-fail-fast |
|
PR_Github #59061 [ run ] triggered by Bot. Commit: |
|
/bot kill |
|
Waiting for #16338 get merged |
|
PR_Github #59080 [ kill ] triggered by Bot. Commit: |
|
PR_Github #59061 [ run ] completed with state |
|
PR_Github #59080 [ kill ] completed with state |
7a55571 to
7078471
Compare
JunyiXu-nv
changed the title
|
/bot run --disable-fail-fast |
|
PR_Github #59353 [ run ] triggered by Bot. Commit: |
7078471 to
f5b352c
Compare
|
PR_Github #59353 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59590 [ run ] triggered by Bot. Commit: |
c4e8f5e to
0170595
Compare
|
PR_Github #59590 [ run ] completed with state
|
…is dead PR NVIDIA#16338 lets the proxy detect workers killed abruptly (MPI_Abort from the PyExecutor HangDetector, SIGKILL, the OOM killer) via OS process monitoring, and PR NVIDIA#15816 unblocks pending requests with a sticky EngineDeadError. But teardown still hangs: shutdown() blocks forever on f.result() for mpi4py futures that never complete after an abrupt kill, then on the result-dispatcher join (the worker that would send the shutdown sentinel is dead), then on joining the dead MPI pool. In CI the detection therefore only moves the client-side hang from generate() into LLM.__exit__, and hang-killed stages still burn pytest's full --timeout=3600 after the GPUs were freed at +300s (e.g. L0_Test-SBSA-Multi-GPU run 2744 GB300-4_GPUs-PyTorch-Post-Merge-2/3 on a post-NVIDIA#15816 main commit). Once the engine is known dead (_engine_dead, set by all detection paths): - give the futures one short collective grace instead of unbounded per-future result() waits, and skip the ones still pending - bound the result-dispatcher join (daemon thread; leaked, not awaited) - shut the owned MPI session down without waiting on the dead pool Orderly-shutdown semantics are unchanged when the engine is alive. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
Hardware validation (GB200, wedge-injection repro) showed the bounded shutdown alone is not enough: after MPI_Abort kills the worker world, generate() unblocks with EngineDeadError ~5s after the kill and shutdown() completes ~20s in, but the process then cannot EXIT. The MPIPoolExecutor manager thread (non-daemon, Thread-1 _manager_spawn) stays blocked in an MPI call on the aborted intercommunicator forever, and both mpi4py's exit hook (THREADS_QUEUES join) and CPython's threading._shutdown non-daemon join wait on it. Add MpiSession.abandon(): a non-waiting shutdown that, for MpiPoolSession, also deregisters the wedged manager thread from both exit-join mechanisms (best-effort, version-tolerant; the thread is reaped with the process). The proxy's dead-engine teardown now abandons the owned session instead of shutting it down. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
Hardware validation round 2: abandoning at teardown is too late. When user code returns normally after handling EngineDeadError, CPython's exit sequence joins non-daemon threads (threading._shutdown) BEFORE the atexit/GC teardown that would run proxy.shutdown() -> abandon(), so the process still hangs on the wedged pool manager thread. Move the abandon call into _mark_engine_dead (the sticky point every detection path funnels through): the engine is terminal there, so the pool can be deregistered from exit joins immediately while user code is still running. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
Hardware validation round 3: the detection-time abandon never ran. In the LLM API path the MpiPoolSession is created by the LLM object and passed into the proxy, so _owns_mpi_session is False and the ownership-gated abandon was skipped -- the process still hung to the outer timeout. Split the concern: release_exit_joins() is non-destructive bookkeeping (deregister the wedged manager thread from mpi4py's and CPython's exit joins, mark the pool dead) and is safe for any session, so _mark_engine_dead now calls it unconditionally. Destructive teardown (abandon/shutdown) remains reserved for the owner. A dead-marked MpiPoolSession also forces wait=False on any later shutdown(), so the owning LLM object's blocking shutdown cannot join the dead pool either. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
0170595 to
00d83fe
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #59746 [ run ] triggered by Bot. Commit: |
|
PR_Github #59746 [ run ] completed with state
|
3 participants
Summary by CodeRabbit
Description
Bound every teardown wait in GenerationExecutorProxy.shutdown() when the engine is known dead (_engine_dead, set by all worker-death detection paths):
Give the mpi4py worker futures one short collective 5 s grace (concurrent.futures.wait) instead of unbounded per-future result() calls, and skip futures still pending.
Bound the result-dispatcher join (5 s): the worker that would send the shutdown sentinel is dead, so the dispatcher may be blocked in a ZMQ recv that never returns. It is a daemon thread — leak it rather than hang teardown.
Shut the owned MPI session down without waiting on the dead pool (shutdown(wait=False)).
Orderly shutdown (engine alive) is byte-for-byte unchanged: futures are reaped blocking, the dispatcher join is unbounded, the session join waits.
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.