[https://nvbugs/6480574][fix] mpi_session: guard MpiPoolSession.shutdown against partial init by lowsfer · Pull Request #16613 · NVIDIA/TensorRT-LLM

lowsfer · 2026-07-20T09:28:56Z

Summary

MpiPoolSession.shutdown() reads self.n_workers (in a log line) and self._wait_shutdown unguarded, but it is reachable on a session that never completed __init__:

from __del__ → shutdown_abort() → shutdown() on a partially-constructed object, and
from the released-session path, e.g. tests/unittest/executor/test_proxy_fast_death.py::test_pool_session_shutdown_never_blocks_after_release, which builds the object with MpiPoolSession.__new__(MpiPoolSession).

In those states the attributes are absent and shutdown() raises AttributeError: 'MpiPoolSession' object has no attribute 'n_workers' (then ... '_wait_shutdown'). The _pool_dead check immediately above already uses getattr(self, '_pool_dead', False) defensively for exactly this reason.

Fix

Guard both accesses with getattr(...), matching the existing _pool_dead style, so a released or partially-constructed session shuts down cleanly instead of raising.

Why main is currently red

This is a cross-PR merge skew, each half green on its own:

the unguarded self.n_workers log line was added by [None][test] disagg startup: detect and fail fast on worker OOM/crash #16456 (2026-07-16), and
the __new__-based test that trips it was added by [TRTLLM-13409][fix] make proxy shutdown non-blocking when the engine is dead #16312 (2026-07-20).

Neither failed in isolation; the combination on main fails deterministically. main has been red on test_pool_session_shutdown_never_blocks_after_release since #16312 landed, and the fix is a small robustness guard.

Test

test_proxy_fast_death.py now passes in full (23 passed, was 1 failed):

python -m pytest tests/unittest/executor/test_proxy_fast_death.py -q
# 23 passed

(No behavior change for fully-initialized sessions: n_workers / _wait_shutdown are always set by __init__; the guards only affect objects that bypass it.)

🤖 Generated with Claude Code

Summary by CodeRabbit

Documentation
- Improved formatting and organization of API documentation.
- Clarified session initialization and external communication examples.
Bug Fixes
- Improved shutdown logging when the worker count is unavailable.
- Preserved existing worker cleanup and shutdown behavior.

coderabbitai · 2026-07-20T09:30:13Z

📝 Walkthrough

Walkthrough

Changes

MPI session updates

Layer / File(s)	Summary
Docstrings and shutdown logging `tensorrt_llm/llmapi/mpi_session.py`	Docstrings were reformatted, and shutdown logging now uses a fallback worker count without changing pool teardown behavior.

Estimated code review effort: 1 (Trivial) | ~2 minutes

Suggested reviewers: cascade812

🚥 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	It matches the required ticket/type format and clearly describes the main fix.
Description check	✅ Passed	It covers the issue, fix, rationale, and test results, with only minor template-heading differences.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

Warning

Review ran into problems

🔥 Problems

Git: Failed to clone repository. Please run the @coderabbitai full review command to re-trigger a full review. If the issue persists, set path_filters to include or exclude specific files.

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

lowsfer · 2026-07-20T10:05:54Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-20T10:11:18Z

PR_Github #60349 [ run ] triggered by Bot. Commit: f036089 Link to invocation

tensorrt-cicd · 2026-07-20T13:26:07Z

PR_Github #60349 [ run ] completed with state SUCCESS. Commit: f036089
/LLM/main/L0_MergeRequest_PR pipeline #48692 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

lowsfer · 2026-07-20T13:41:03Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-20T13:46:54Z

PR_Github #60380 [ run ] triggered by Bot. Commit: f036089 Link to invocation

tensorrt-cicd · 2026-07-20T16:39:34Z

PR_Github #60380 [ run ] completed with state FAILURE. Commit: f036089
/LLM/main/L0_MergeRequest_PR pipeline #48719 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

lowsfer · 2026-07-20T17:34:44Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-20T17:40:33Z

PR_Github #60420 [ run ] triggered by Bot. Commit: f036089 Link to invocation

tensorrt-cicd · 2026-07-20T19:19:33Z

PR_Github #60420 [ run ] completed with state FAILURE. Commit: f036089
/LLM/main/L0_MergeRequest_PR pipeline #48757 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

lowsfer · 2026-07-21T05:50:59Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-21T05:56:35Z

PR_Github #60600 [ run ] triggered by Bot. Commit: f036089 Link to invocation

tensorrt-cicd · 2026-07-21T07:03:31Z

PR_Github #60600 [ run ] completed with state SUCCESS. Commit: f036089
/LLM/main/L0_MergeRequest_PR pipeline #48908 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

longlee0622 · 2026-07-21T07:14:21Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-21T07:22:14Z

PR_Github #60624 [ run ] triggered by Bot. Commit: f036089 Link to invocation

…own against partial init MpiPoolSession.shutdown() reads self.n_workers (in a log line) and self._wait_shutdown unguarded, but shutdown() is reachable on a session that never completed __init__: - from __del__ -> shutdown_abort() -> shutdown() on a partially-constructed object, and - from the released-session path, e.g. tests/unittest/executor/test_proxy_fast_death.py::test_pool_session_shutdown_never_blocks_after_release, which builds the object with MpiPoolSession.__new__(MpiPoolSession). In those states the attributes are absent and shutdown() raises AttributeError: 'MpiPoolSession' object has no attribute 'n_workers' (then '_wait_shutdown'). Guard both accesses with getattr(...), matching the existing _pool_dead guard, so a released or partially-constructed session shuts down cleanly instead of raising. No behavior change for fully-initialized sessions (n_workers / _wait_shutdown are always set by __init__; the guards only affect objects that bypass it). Also remove the temporary waiver for that test (nvbugs/6480574) so CI re-validates the now-fixed test. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

lowsfer · 2026-07-21T08:04:29Z

Rebased onto latest main and removed the temporary waiver for test_proxy_fast_death.py::test_pool_session_shutdown_never_blocks_after_release (nvbugs/6480574, added by the auto-waive PR #16615) so CI now re-validates the previously-failing test.

This PR is the fix for nvbugs/6480574 / TRTLLM-14480. Empirical confirmation of the fix from recent CI on the same test:

Build L0_MergeRequest_PR #48602 (unguarded shutdown()) → AttributeError: 'MpiPoolSession' object has no attribute 'n_workers'
The next commit adding the getattr guard passed with the same test present (builds #48694 / #48718)

The failure is a cross-PR merge skew (unguarded n_workers log line from #16456 + the __new__-based test from #16312); deterministic on main, hence the auto-waive.

lowsfer · 2026-07-21T08:04:59Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-21T08:10:20Z

PR_Github #60644 [ run ] triggered by Bot. Commit: 299930b Link to invocation

tensorrt-cicd · 2026-07-21T08:13:55Z

PR_Github #60624 [ run ] completed with state ABORTED. Commit: f036089

Link to invocation

lowsfer · 2026-07-21T09:23:49Z

Closing in favor of #16630 by @JunyiXu-nv (the nvbugs/6480574 assignee), which fixes the same failure test-side by seeding the __new__-bypassed attributes in test_pool_session_shutdown_never_blocks_after_release.

For the record, empirical CI confirmation of the root cause (same test): build L0_MergeRequest_PR #48602 (unguarded) failed with AttributeError: 'MpiPoolSession' object has no attribute 'n_workers'; the next commit adding the guard passed (#48694/#48718). Thanks @JunyiXu-nv for taking it.

tensorrt-cicd · 2026-07-21T11:32:13Z

PR_Github #60644 [ run ] completed with state FAILURE. Commit: 299930b
/LLM/main/L0_MergeRequest_PR pipeline #48947 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

lowsfer requested a review from a team as a code owner July 20, 2026 09:28

lowsfer requested review from Tabrizian and cascade812 July 20, 2026 09:28

github-actions Bot assigned lowsfer Jul 20, 2026

coderabbitai Bot mentioned this pull request Jul 20, 2026

[https://nvbugs/6480574][fix] seed attrs bypassed by __new__ in pool session shutdown test #16630

Merged

1 task

mikeiovine removed the request for review from Tabrizian July 20, 2026 15:43

lowsfer force-pushed the fix/mpi-pool-session-shutdown-partial-init branch from f036089 to 299930b Compare July 21, 2026 08:04

lowsfer requested review from a team as code owners July 21, 2026 08:04

lowsfer requested review from brnguyen2, mzweilz, schetlur-nv and xinhe-nv July 21, 2026 08:04

lowsfer requested a review from mlefeb01 July 21, 2026 08:04

lowsfer changed the title ~~[None][fix] mpi_session: guard MpiPoolSession.shutdown against partial init~~ [https://nvbugs/6480574][fix] mpi_session: guard MpiPoolSession.shutdown against partial init Jul 21, 2026

lowsfer closed this Jul 21, 2026

Uh oh!

Conversation

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

Uh oh!

Summary

Fix

Why main is currently red

Test

Summary by CodeRabbit

Uh oh!

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

Uh oh!

Walkthrough

Changes

Review ran into problems

Uh oh!

lowsfer commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

lowsfer commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

lowsfer commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

lowsfer commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

longlee0622 commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

lowsfer commented Jul 21, 2026

Uh oh!

lowsfer commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

lowsfer commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

lowsfer commented Jul 20, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 20, 2026 •

edited

Loading