[TRTLLM-13409][fix] make proxy shutdown non-blocking when the engine is dead by JunyiXu-nv · Pull Request #16312 · NVIDIA/TensorRT-LLM

JunyiXu-nv · 2026-07-13T08:05:07Z

Summary by CodeRabbit

Bug Fixes
- Improved detection of abruptly terminated worker processes, including crashes and forced kills.
- Pending requests now receive engine-dead errors more promptly when a worker fails.
- Prevented executor shutdown from hanging after an engine failure.
- Improved startup error reporting when worker initialization fails.
- Added safeguards to distinguish restarted processes from stale worker identities.

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-compatible or api-breaking. For api-breaking, include BREAKING in 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.

coderabbitai · 2026-07-13T08:11:31Z

📝 Walkthrough

Walkthrough

The proxy now receives worker PID metadata, monitors local worker liveness with psutil, propagates abrupt worker death through the existing fast-death path, and bounds shutdown waits when the engine is already dead.

Changes

Worker liveness monitoring

Layer / File(s)	Summary
Process identity handshake `tensorrt_llm/executor/worker.py`, `tensorrt_llm/executor/proxy.py`	Workers gather rank, hostname, PID, and creation time, then return the metadata through the startup handshake for local registration.
PID liveness detection `tensorrt_llm/executor/proxy.py`, `tests/unittest/executor/test_proxy_fast_death.py`	The proxy validates local worker PIDs, detects zombie and recycled processes, integrates polling into health and error monitoring, and tests death propagation and registration filtering.
Dead-engine shutdown `tensorrt_llm/executor/proxy.py`	Shutdown uses bounded waits and joins, skips incomplete future results, and avoids waiting for MPI teardown when the engine is dead.

Estimated code review effort: 4 (Complex) | ~45 minutes

Possibly related PRs

NVIDIA/TensorRT-LLM#16053: Modifies overlapping GenerationExecutorProxy shutdown and MPI session teardown paths.

Suggested reviewers: schetlur-nv

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title follows the required ticket/type pattern and clearly summarizes the main change.
Description check	✅ Passed	The required sections are present and the change is clearly explained, but the Test Coverage section is left blank.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

coderabbitai

🧹 Nitpick comments (1)

tests/unittest/executor/test_proxy_fast_death.py (1)
392-456: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Coverage 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, None no-op) well.

Gap: the dead-engine bounded-shutdown behavior added in tensorrt_llm/executor/proxy.py (the _engine_dead-gated concurrent.futures.wait(..., timeout=5.0), bounded dispatch_result_thread.join, and mpi_session.shutdown(wait=not self._engine_dead)) has no unit coverage. Consider a test with a fake session/futures asserting that shutdown() returns within the bound and does not block on pending futures when _engine_dead is 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

📥 Commits

Reviewing files that changed from the base of the PR and between 53d2bde and 7a55571.

📒 Files selected for processing (3)

tensorrt_llm/executor/proxy.py
tensorrt_llm/executor/worker.py
tests/unittest/executor/test_proxy_fast_death.py

QiJune

LGTM

JunyiXu-nv · 2026-07-14T01:35:30Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-14T01:42:42Z

PR_Github #59061 [ run ] triggered by Bot. Commit: 7a55571 Link to invocation

JunyiXu-nv · 2026-07-14T02:54:47Z

/bot kill

JunyiXu-nv · 2026-07-14T02:55:08Z

Waiting for #16338 get merged

tensorrt-cicd · 2026-07-14T03:01:41Z

PR_Github #59080 [ kill ] triggered by Bot. Commit: 7a55571 Link to invocation

tensorrt-cicd · 2026-07-14T03:05:53Z

PR_Github #59061 [ run ] completed with state ABORTED. Commit: 7a55571

Link to invocation

tensorrt-cicd · 2026-07-14T03:06:03Z

PR_Github #59080 [ kill ] completed with state SUCCESS. Commit: 7a55571
Successfully killed previous jobs for commit 7a55571

Link to invocation

JunyiXu-nv · 2026-07-15T03:36:54Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-15T03:42:05Z

PR_Github #59353 [ run ] triggered by Bot. Commit: 7078471 Link to invocation

tensorrt-cicd · 2026-07-15T13:23:18Z

PR_Github #59353 [ run ] completed with state SUCCESS. Commit: 7078471
/LLM/main/L0_MergeRequest_PR pipeline #47830 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

JunyiXu-nv · 2026-07-16T02:49:40Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-16T02:55:18Z

PR_Github #59590 [ run ] triggered by Bot. Commit: c4e8f5e Link to invocation

tensorrt-cicd · 2026-07-16T15:55:26Z

PR_Github #59590 [ run ] completed with state SUCCESS. Commit: c4e8f5e
/LLM/main/L0_MergeRequest_PR pipeline #48033 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

…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>

JunyiXu-nv · 2026-07-16T16:09:22Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-16T16:14:54Z

PR_Github #59746 [ run ] triggered by Bot. Commit: 00d83fe Link to invocation

tensorrt-cicd · 2026-07-16T19:55:32Z

PR_Github #59746 [ run ] completed with state SUCCESS. Commit: 00d83fe
/LLM/main/L0_MergeRequest_PR pipeline #48172 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

JunyiXu-nv · 2026-07-17T08:07:17Z

/bot run

tensorrt-cicd · 2026-07-17T08:13:48Z

PR_Github #59936 [ run ] triggered by Bot. Commit: 00d83fe Link to invocation

tensorrt-cicd · 2026-07-17T12:12:58Z

PR_Github #59936 [ run ] completed with state SUCCESS. Commit: 00d83fe
/LLM/main/L0_MergeRequest_PR pipeline #48336 completed with status: 'SUCCESS'

CI Report

Link to invocation

…l init MpiPoolSession.shutdown() is reachable from __del__ (via shutdown_abort) and may run on a session that never completed __init__ -- e.g. after release_exit_joins() marks the pool dead, or in the released-session path exercised by test_proxy_fast_death.py:: test_pool_session_shutdown_never_blocks_after_release, which builds the object with __new__. shutdown() read self.n_workers (log line) and self._wait_shutdown unguarded, raising AttributeError, while the _pool_dead check right above already used getattr defensively. Guard both accesses with getattr, matching the existing _pool_dead style, so a released/partially-constructed session shuts down without crashing. This unbreaks a cross-PR merge-skew regression on main: the unguarded self.n_workers log came from NVIDIA#16456 while the __new__-based test came from NVIDIA#16312; each was green alone but the combination fails deterministically. Also add the missing D205 docstring blank lines the ruff-legacy gate flagged on the touched file. Verified: the whole test_proxy_fast_death.py file now passes (23 passed, was 1 failed). Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

…ool-session test test_proxy_fast_death builds GenerationExecutorProxy / MpiPoolSession via __new__ (bypassing __init__), then drives shutdown(). A recent main change (NVIDIA#16523, multi-process HTTP frontends) made GenerationExecutorProxy.shutdown() -> _cleanup_multi_frontend_ipc_dir() read self._multi_frontend_ipc_dir unguarded, so the __new__-built proxies raise AttributeError at GC/teardown (test_shutdown_does_not_block_on_dead_engine and the two sibling shutdown tests). This is a cross-PR merge-skew: NVIDIA#16312 added the __new__-based tests, NVIDIA#16523 later added the unguarded attribute access; neither PR tested the combination. Seed _multi_frontend_ipc_dir / _multi_frontend_hmac in the shared _bare_proxy() helper so the teardown path is a clean no-op, mirroring how NVIDIA#16630 (nvbugs/6480574) handled the earlier n_workers/_wait_shutdown variant. Also re-align with main: restore the n_workers/_wait_shutdown seeding in test_pool_session_shutdown_never_blocks_after_release and drop its now-removed waive (main dropped it in NVIDIA#16630). Whole file: 23 passed (was 3 failed + 1 waived/broken). Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

…roxy_fast_death __new__ seeding Two base-branch (main) regressions surfaced by rebasing onto latest main; neither is in the KVCacheManagerV2 change set. 1) tensorrt_llm import break (nvbugs/6442074): SpecWorkerBase.__init_subclass__ forbids subclasses from overriding forward() -- they must implement _forward_impl so the base forward() wrapper can guarantee spec-dec attn-metadata cleanup. MTPEagleDynamicTreeWorker still overrode forward(), so 'import tensorrt_llm' raised TypeError at class-definition time via the normal _torch.models -> speculative -> mtp_dynamic_tree import chain, breaking all test collection / package sanity. Rename its forward() to _forward_impl(), matching every sibling worker (MTPEagleWorker, DSparkWorker, eagle3, eagle3_dynamic_tree); callers use the base forward() wrapper so no caller changes. Verified: 'import tensorrt_llm' succeeds again. 2) test_proxy_fast_death cross-PR merge-skew: NVIDIA#16312 added __new__-built proxy tests; NVIDIA#16523 later made GenerationExecutorProxy.shutdown() -> _cleanup_multi_frontend_ipc_dir() read self._multi_frontend_ipc_dir unguarded, so the __new__ proxies raised AttributeError at teardown. Seed _multi_frontend_ipc_dir/_multi_frontend_hmac in the shared _bare_proxy() helper so GC-time teardown is a clean no-op, mirroring NVIDIA#16630 (nvbugs/6480574) for the earlier n_workers variant. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

JunyiXu-nv requested a review from a team as a code owner July 13, 2026 08:05

JunyiXu-nv requested a review from schetlur-nv July 13, 2026 08:05

github-actions Bot assigned JunyiXu-nv Jul 13, 2026

coderabbitai Bot reviewed Jul 13, 2026

View reviewed changes

JunyiXu-nv requested review from QiJune and removed request for schetlur-nv July 13, 2026 08:48

QiJune approved these changes Jul 13, 2026

View reviewed changes

JunyiXu-nv force-pushed the dev-junyix-fix-proxy-worker-liveness branch from 7a55571 to 7078471 Compare July 15, 2026 03:17

JunyiXu-nv requested a review from a team as a code owner July 15, 2026 03:17

JunyiXu-nv requested review from asfiyab-nvidia and tongyuantongyu July 15, 2026 03:17

JunyiXu-nv changed the title ~~[TRTLLM-13409][fix] detect worker death invisible to MPI futures via PID liveness~~ [TRTLLM-13409][fix] make proxy shutdown non-blocking when the engine is dead Jul 15, 2026

JunyiXu-nv force-pushed the dev-junyix-fix-proxy-worker-liveness branch from 7078471 to f5b352c Compare July 15, 2026 05:32

JunyiXu-nv force-pushed the dev-junyix-fix-proxy-worker-liveness branch from c4e8f5e to 0170595 Compare July 16, 2026 14:25

JunyiXu-nv added 2 commits July 17, 2026 00:08

JunyiXu-nv added 3 commits July 17, 2026 00:08

[TRTLLM-13409][chore] tighten comments to constraint-only form

00d83fe

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>

JunyiXu-nv force-pushed the dev-junyix-fix-proxy-worker-liveness branch from 0170595 to 00d83fe Compare July 16, 2026 16:08

JunyiXu-nv merged commit 5a4d258 into NVIDIA:main Jul 20, 2026
7 checks passed

lowsfer mentioned this pull request Jul 20, 2026

[https://nvbugs/6480574][fix] mpi_session: guard MpiPoolSession.shutdown against partial init #16613

Closed

Uh oh!

Conversation

JunyiXu-nv commented Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

coderabbitai Bot commented Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

QiJune left a comment

Choose a reason for hiding this comment

Uh oh!

JunyiXu-nv commented Jul 14, 2026

Uh oh!

tensorrt-cicd commented Jul 14, 2026

Uh oh!

JunyiXu-nv commented Jul 14, 2026

Uh oh!

JunyiXu-nv commented Jul 14, 2026

Uh oh!

tensorrt-cicd commented Jul 14, 2026

Uh oh!

tensorrt-cicd commented Jul 14, 2026

Uh oh!

tensorrt-cicd commented Jul 14, 2026

Uh oh!

JunyiXu-nv commented Jul 15, 2026

Uh oh!

tensorrt-cicd commented Jul 15, 2026

Uh oh!

tensorrt-cicd commented Jul 15, 2026

Uh oh!

JunyiXu-nv commented Jul 16, 2026

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

JunyiXu-nv commented Jul 16, 2026

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

JunyiXu-nv commented Jul 17, 2026

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

JunyiXu-nv commented Jul 13, 2026 •

edited

Loading

coderabbitai Bot commented Jul 13, 2026 •

edited

Loading