[https://nvbugs/6581066][fix] In that branch, call shutdown_abort() when the session is owned, then… - #17511
Closed
trtllm-agent wants to merge 1 commit into
Closed
[https://nvbugs/6581066][fix] In that branch, call shutdown_abort() when the session is owned, then…#17511trtllm-agent wants to merge 1 commit into
trtllm-agent wants to merge 1 commit into
Conversation
…during init A rank that fails executor setup only notifies the proxy if it is the leader; a non-leader just returns. Its peers therefore stay blocked in the init collective still holding their share of the weights, and the proxy's init wait loop raised without marking the engine dead or aborting anyone. The surviving ranks then leaked until job end and the LLM's own blocking mpi_session.shutdown() joined a dead world instead of reporting the failure, turning a fast init error into a test timeout. Do what the adjacent init-error branch already does: abort a session we own, then record the death. Order matters -- _mark_engine_dead() reaches release_exit_joins(), which marks the pool dead and forces shutdown() non-blocking, so marking first would defang shutdown_abort()'s MPI_Abort escalation. Recording the death also releases exit joins for a borrowed session, so its owner's later join cannot block forever. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
Contributor
WalkthroughThe proxy now centralizes MPI session abortion for owned sessions. Worker death during initialization marks the engine dead, releases pending requests, and triggers shutdown handling. Tests cover owned and externally owned sessions and verify operation ordering. ChangesProxy initialization failure handling
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Worker
participant Proxy
participant MPI
participant PendingResults
Worker->>Proxy: Fail during initialization
Proxy->>MPI: shutdown_abort() for owned session
Proxy->>PendingResults: Broadcast EngineDeadError
Proxy->>Proxy: Mark engine dead and release exit joins
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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.
Inline comments:
In `@tensorrt_llm/executor/proxy.py`:
- Around line 668-683: Update tensorrt_llm/executor/proxy.py:668-683 and 686-691
so every initialization failure, including exceptional mpi_done_callback futures
and non-ready status errors, calls _abort_owned_session() before
_mark_engine_dead(), then raises; defer callback bookkeeping during
initialization if needed to preserve this ordering. Add coverage in
tests/unittest/executor/test_proxy_fast_death.py:560-629 for exceptional futures
and both owned and borrowed status-error paths, while retaining the existing
successful-future test.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d2c34164-cf7c-429b-a0e7-378f5fac30df
📒 Files selected for processing (2)
tensorrt_llm/executor/proxy.pytests/unittest/executor/test_proxy_fast_death.py
Comment on lines
+668
to
+683
| death = RuntimeError( | ||
| "Executor worker died during initialization") | ||
| # A non-leader rank that fails here returns without notifying | ||
| # anyone, so its peers stay blocked in the init collective | ||
| # still holding their share of the weights. Raising alone | ||
| # leaks them until job end and makes the next blocking | ||
| # shutdown() hang instead of reporting this failure. | ||
| # | ||
| # Abort before marking: shutdown_abort() escalates to | ||
| # MPI_Abort only if its own blocking shutdown() overruns the | ||
| # grace period, and _mark_engine_dead() -> release_exit_joins() | ||
| # marks the pool dead, which forces that shutdown() | ||
| # non-blocking. Marking first would defang the abort. | ||
| self._abort_owned_session(death) | ||
| self._mark_engine_dead(death) | ||
| raise death |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- proxy outline ---'
ast-grep outline tensorrt_llm/executor/proxy.py | sed -n '1,220p'
printf '%s\n' '--- relevant proxy symbols ---'
rg -n -C 8 '_start_executor_workers|_handle_worker_death|mpi_done_callback|_abort_owned_session|_mark_engine_dead|shutdown_abort|worker_init_status_queue' tensorrt_llm/executor/proxy.py
printf '%s\n' '--- relevant tests ---'
rg -n -C 8 '_proxy_awaiting_worker_init|worker_death_during_init|set_exception|init.*error|shutdown_abort|release_exit_joins' tests/unittest/executor/test_proxy_fast_death.py
printf '%s\n' '--- repository test/config references ---'
rg -n 'test_proxy_fast_death|pytest tests/unittest|test-db|coverage verdict' tests qa scripts 2>/dev/null | sed -n '1,220p'Repository: NVIDIA/TensorRT-LLM
Length of output: 42888
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- proxy lifecycle and shutdown ---'
sed -n '300,365p;615,705p;740,845p' tensorrt_llm/executor/proxy.py
printf '%s\n' '--- session implementations ---'
fd -i -t f 'mpi*session*' tensorrt_llm
rg -n -C 12 'def (shutdown_abort|shutdown|release_exit_joins)|release_exit_joins|shutdown_abort' tensorrt_llm tests/unittest/executor/test_proxy_fast_death.py | sed -n '1,320p'
printf '%s\n' '--- worker initialization status producer ---'
rg -n -C 12 'ready_signal|worker_init_status_queue|error_trace|READY_SIGNAL' tensorrt_llm/executor tensorrt_llm | sed -n '1,360p'
printf '%s\n' '--- test registration ---'
sed -n '110,130p' tests/integration/test_lists/test-db/l0_a10.yml
printf '%s\n' '--- test file header and relevant setup ---'
sed -n '1,80p;330,570p' tests/unittest/executor/test_proxy_fast_death.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- session definitions ---'
rg -l '^class (MpiPoolSession|MpiCommSession|RemoteMpiCommSessionClient)|def shutdown_abort|def release_exit_joins' tensorrt_llm | sort
printf '%s\n' '--- exact session methods ---'
rg -n -C 20 '^ def (shutdown_abort|shutdown|release_exit_joins|abandon)' tensorrt_llm/llmapi tensorrt_llm | head -n 420
printf '%s\n' '--- worker_main definitions and status sends ---'
rg -n -C 20 'def worker_main|worker_init_status_queue.*put|init.*status|error_trace|ready_signal' tensorrt_llm/executor tensorrt_llm | head -n 500
printf '%s\n' '--- focused test helpers and mocks ---'
sed -n '1,45p;320,370p;543,635p' tests/unittest/executor/test_proxy_fast_death.py
printf '%s\n' '--- tracked changes summary ---'
git diff --statRepository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- all ready_signal producers ---'
rg -n -C 10 --glob '*.py' 'ready_signal|worker_init_status_queue' . | grep -v '^./tests/' | sed -n '1,360p'
printf '%s\n' '--- focused proxy and session source ---'
sed -n '107,175p;300,340p;423,430p' tensorrt_llm/llmapi/mpi_session.py
sed -n '615,695p' tensorrt_llm/executor/proxy.py
printf '%s\n' '--- static control-flow verifier ---'
python3 - <<'PY'
from pathlib import Path
proxy = Path("tensorrt_llm/executor/proxy.py").read_text()
session = Path("tensorrt_llm/llmapi/mpi_session.py").read_text()
callback = proxy[proxy.index("def mpi_done_callback"):proxy.index("tracer_init_kwargs")]
init = proxy[proxy.index("while True:", proxy.index("def _start_executor_workers")):
proxy.index("def _register_worker_processes")]
shutdown_abort = session[session.index("def shutdown_abort"):session.index("def release_exit_joins")]
release = session[session.index("def release_exit_joins"):session.index("def abort", session.index("def release_exit_joins"))]
pool_shutdown = session[session.index("def shutdown(self, wait=True):", session.index("class MpiPoolSession")):
session.index("def _collect_worker_identities")]
checks = {
"callback_marks_engine_dead": "_handle_worker_death" in callback,
"init_branch_aborts_before_marks": (
"self._abort_owned_session(death)" in init
and init.index("self._abort_owned_session(death)")
< init.index("self._mark_engine_dead(death)")
),
"status_branch_aborts_without_mark": (
"self._abort_owned_session(ready_signal)" in init
and "self._mark_engine_dead(ready_signal)" not in init
),
"release_marks_pool_dead": "self._pool_dead = True" in release,
"dead_pool_forces_nonblocking_shutdown": (
"if getattr(self, '_pool_dead', False)" in pool_shutdown
and "wait = False" in pool_shutdown
),
"shutdown_abort_waits_then_timeout_abort": (
"self.shutdown()" in shutdown_abort
and "self.abort()" not in shutdown_abort
and "self._abort_on_timeout" in shutdown_abort
),
}
for name, value in checks.items():
print(f"{name}={value}")
if not all(checks.values()):
raise SystemExit("control-flow invariant failed")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 26985
Finalize all initialization failures before marking the engine dead.
- An exceptional
mpi_done_callbackcan call_mark_engine_dead()before the initialization loop.release_exit_joins()then forcesMpiPoolSession.shutdown()to usewait=False, which can preventshutdown_abort()from reachingMPI_Abort. Defer callback bookkeeping during initialization or preserve abort-before-mark ordering. - The non-ready status path must call
_mark_engine_dead()after_abort_owned_session()and before raising. Otherwise shutdown can wait on unfinished futures and borrowed sessions do not release exit joins. - Add tests for exceptional futures and owned and borrowed status-error paths. The existing future test completes successfully, and
tests/unittest/executor/test_proxy_fast_death.pyis already registered intests/integration/test_lists/test-db/l0_a10.yml.
📍 Affects 2 files
tensorrt_llm/executor/proxy.py#L668-L683(this comment)tensorrt_llm/executor/proxy.py#L686-L691tests/unittest/executor/test_proxy_fast_death.py#L560-L629
🤖 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 `@tensorrt_llm/executor/proxy.py` around lines 668 - 683, Update
tensorrt_llm/executor/proxy.py:668-683 and 686-691 so every initialization
failure, including exceptional mpi_done_callback futures and non-ready status
errors, calls _abort_owned_session() before _mark_engine_dead(), then raises;
defer callback bookkeeping during initialization if needed to preserve this
ordering. Add coverage in
tests/unittest/executor/test_proxy_fast_death.py:560-629 for exceptional futures
and both owned and borrowed status-error paths, while retaining the existing
successful-future test.
Sources: Coding guidelines, Path instructions
BowenFu
requested changes
Aug 13, 2026
BowenFu
left a comment
BowenFu
left a comment
Contributor
There was a problem hiding this comment.
Changes required: preserve abort-before-marking on every initialization failure.
- The exceptional-future path violates the claimed invariant:
add_done_callback()calls_handle_worker_death(), which calls_mark_engine_dead()before the polling loop reaches_abort_owned_session(). For an already-failed future, the callback runs synchronously during registration. This can makeshutdown_abort()non-blocking and leave surviving ranks alive. - The non-ready status path aborts but never calls
_mark_engine_dead(), so pending-result and exit-join cleanup is inconsistent.
Please defer initialization-time callback bookkeeping or route all three paths through one ordered helper: abort the owned session, then mark dead, then raise. Add focused owned/borrowed regressions for an exceptional future and non-ready status, while retaining successful-future coverage. This is required for this PR.
zhaoyangwang-nvidia
requested changes
Aug 14, 2026
| # (shared) session must stay alive for its owner to tear down. | ||
| if self._owns_mpi_session: | ||
| self.mpi_session.shutdown_abort(reason=ready_signal) | ||
| self._abort_owned_session(ready_signal) |
Collaborator
There was a problem hiding this comment.
This branch is the same failure as the loop above ("worker world unusable after init") but only aborts, leaving _engine_dead False. A later shutdown() then skips the 5s future grace, calls f.result() with no timeout on futures of ranks still wedged in the init collective, and takes the blocking mpi_session.shutdown() path instead of abandon(). For a borrowed session (_owns_mpi_session=False) nothing happens at all here, so the hang described in the PR body remains on this path. Should this call _mark_engine_dead(ready_signal) after _abort_owned_session(...), for symmetry with the loop above?
Merged
2 tasks
Collaborator
|
Superseded by #17817, which carries this commit and additionally addresses the review feedback: the fast-death callback registration is deferred until the worker world reported ready (an already-failed future ran |
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
4 participants
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.
Summary
pytest tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_8gpus[attention_dp_off-trtllm] -vTest plan
Links
Dev Engineer Review
_mark_engine_dead()._mark_engine_dead()makes shutdown non-blocking.QA Engineer Review
tests/unittest/executor/test_proxy_fast_death.pyfor:shutdown_abort().EngineDeadErrordelivery to pending results.shutdown_abort()ordering before engine-dead bookkeeping.tests/integration/test_lists/entries were changed.