[https://nvbugs/6435642][fix] handle session reuse worker registration

[chienchunhung]

@coderabbitai summary

Description

PR #16338 added MPI worker-process monitoring for NVBug 6435642, but its startup path used isinstance(self.mpi_session, MpiPoolSession). The repository-wide session-reuse test hook replaces the module-local MpiPoolSession class reference with a factory function, causing isinstance() to raise TypeError under import-order-dependent test execution.

Identify locally spawned pool sessions by excluding MpiCommSession and RemoteMpiCommSessionClient, matching the existing session classification used earlier in GenerationExecutorProxy. Extract worker identity registration into a small helper and add deterministic regression coverage that explicitly replaces proxy.MpiPoolSession with a factory before exercising the production registration logic.

Remove the two out-of-tree test waivers introduced by #16423 so CI executes the affected session-reuse coverage. This supersedes the bot-owned fix in #16417 and reverses the temporary containment introduced by #16423. The unrelated DeepSeekV3Lite waiver remains unchanged.

Related work:

• [https://nvbugs/6435642][fix] detect killed MPI executor workers #16338
• [https://nvbugs/6426847][fix] Invert the check to `not isinstance(self.mpi_session, (MpiCommSession… #16417
• [None][infra] Waive 2 failed cases for main in pre-merge 47854 #16423
• https://nvbugs/6435642

Test Coverage

Targeted CI passed on current commit a9d549c1f4be02fbbd9090a8535dc1bf806e05d6:

• Command: /bot run --disable-fail-fast --stage-list "A30-PyTorch-1, A30-PyTorch-2, GB10-PyTorch-1, RTXPro6000D-PyTorch-1"
• Pipeline: L0_MergeRequest_PR #47947 — SUCCESS (Partly Tested)
• Results: 687 passed, 0 failed, 397 skipped
• All requested stages passed.

The newly unwaived tests were collected and passed with no skip reason:

Test	A30 (A30-PyTorch-2)	GB10	RTX Pro 6000D
unittest/_torch/modeling -k "modeling_out_of_tree"	Passed	Passed	Passed
TestOutOfTree::test_llm_api[True]	Passed	Passed	Passed

Regression verification: These are the same tests that previously failed after #16338 and were temporarily waived by #16423. With both waivers removed in this PR, the tests are now passing on A30, GB10, and RTX Pro 6000D.

Local validation:

• Focused pre-commit hooks passed for all three changed files.
• Python syntax compilation passed.
• Waiver duplicate and AST validation passed.
• git diff --check passed.

Remaining verification:

• The new deterministic unit test was not selected by targeted pipeline #47947.
• A full /bot run --disable-fail-fast has been requested on the same commit and is awaiting bot acknowledgement/results.

PR Checklist

• PR description clearly explains what and why.
• PR follows the TensorRT-LLM coding guidelines to the best of my knowledge.
• Deterministic regression coverage is included for the fixed path.
• The previously failing out-of-tree tests passed unwaived on A30, GB10, and RTX Pro 6000D.
• No public API or dependency changes are introduced.
• Full current-head bot CI passes.

[chienchunhung]

/bot run --disable-fail-fast --stage-list "A30-PyTorch-1, A30-PyTorch-2, GB10-PyTorch-1, RTXPro6000D-PyTorch-1"

[tensorrt-cicd]

PR_Github #59488 [ run ] triggered by Bot. Commit: a9d549c Link to invocation

[tensorrt-cicd]

PR_Github #59488 [ run ] completed with state SUCCESS. Commit: a9d549c
/LLM/main/L0_MergeRequest_PR pipeline #47947 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

[chienchunhung]

/bot run --disable-fail-fast

[coderabbitai]

📝 Walkthrough

Walkthrough

GenerationExecutorProxy now centralizes worker identity registration and supports reused MPI sessions by excluding external communication sessions. A unit test covers the factory-based session path, and obsolete modeling waivers are removed.

Changes

Session-Reuse Worker Registration

Layer / File(s)	Summary
Worker identity registration tensorrt_llm/executor/proxy.py	Worker startup delegates identity registration to _register_worker_processes, which registers identities for non-communication sessions when the status payload has three elements.
Registration validation and waiver cleanup tests/unittest/executor/test_proxy_fast_death.py, tests/integration/test_lists/waives.txt	A unit test validates registration with a reused session factory, and out-of-tree modeling waivers are removed while the Qwen waiver remains.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#16338: Updates proxy worker identity registration from the worker initialization payload.
• NVIDIA/TensorRT-LLM#16411: Changes proxy registration gating for reused sessions.
• NVIDIA/TensorRT-LLM#16431: Replaces the MpiPoolSession check with communication-session-based registration logic.

Suggested reviewers: tburt-nv, pcastongauy, niukuo, junyixu-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 matches the main fix and follows the requested [ticket][type] summary format.
Description check	✅ Passed	The description includes the required Description, Test Coverage, and PR Checklist sections with relevant details.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

🧹 Nitpick comments (2)

tensorrt_llm/executor/proxy.py (1)

578-589: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Apply precise Python annotations across the new code.

• tensorrt_llm/executor/proxy.py#L578-L589: replace bare tuple/List[...] with a precise ready-status tuple alias using built-in generics.
• tests/unittest/executor/test_proxy_fast_death.py#L102-L102: add -> None to the new test.

As per coding guidelines, Python functions require precise annotations and built-in generic types are preferred.

🤖 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 578 - 589, Update
tensorrt_llm/executor/proxy.py lines 578-589 in _register_worker_processes to
use the precise ready-status tuple alias with built-in generic syntax instead of
bare tuple and List[...] types. Update
tests/unittest/executor/test_proxy_fast_death.py line 102 to add a -> None
return annotation to the new test.

Source: Coding guidelines

tests/unittest/executor/test_proxy_fast_death.py (1)

102-113: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Cover the excluded session types.

This test verifies only the factory-registration path. Add or confirm tests in tests/unittest/executor/test_proxy_fast_death.py for both MpiCommSession and RemoteMpiCommSessionClient, asserting that register() is not called; otherwise an incorrect classification would still pass.

As per path instructions, coverage is insufficient for the changed classification and needs follow-up.

🤖 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 102 - 113,
Extend test_register_worker_processes_with_session_reuse_factory or add focused
tests in the same test module for MpiCommSession and RemoteMpiCommSessionClient.
For each session type, configure the bare proxy with that session, invoke
_register_worker_processes, and assert _worker_process_monitor.register was not
called, while preserving the existing factory-registration assertion.

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 `@tensorrt_llm/executor/proxy.py`:
- Around line 578-589: Update tensorrt_llm/executor/proxy.py lines 578-589 in
_register_worker_processes to use the precise ready-status tuple alias with
built-in generic syntax instead of bare tuple and List[...] types. Update
tests/unittest/executor/test_proxy_fast_death.py line 102 to add a -> None
return annotation to the new test.

In `@tests/unittest/executor/test_proxy_fast_death.py`:
- Around line 102-113: Extend
test_register_worker_processes_with_session_reuse_factory or add focused tests
in the same test module for MpiCommSession and RemoteMpiCommSessionClient. For
each session type, configure the bare proxy with that session, invoke
_register_worker_processes, and assert _worker_process_monitor.register was not
called, while preserving the existing factory-registration assertion.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 264bad9a-d9c0-4620-9ee6-a66428a8233f

📥 Commits

Reviewing files that changed from the base of the PR and between 0623958 and a9d549c.

📒 Files selected for processing (3)

• tensorrt_llm/executor/proxy.py
• tests/integration/test_lists/waives.txt
• tests/unittest/executor/test_proxy_fast_death.py

💤 Files with no reviewable changes (1)

• tests/integration/test_lists/waives.txt

[tensorrt-cicd]

PR_Github #59523 [ run ] triggered by Bot. Commit: a9d549c Link to invocation

[tensorrt-cicd]

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

/bot run

[tensorrt-cicd]

PR_Github #59683 [ run ] triggered by Bot. Commit: a9d549c Link to invocation

[tensorrt-cicd]

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

[shuyixiong]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #59696 [ run ] triggered by Bot. Commit: a9d549c Link to invocation

[BowenFu]

LGTM — proper fix for MPI session-reuse worker registration: the exclusion-check (not comm/remote session) matches the init classification and avoids dereferencing the swapped MpiPoolSession symbol; len==3 guard preserved.

[litaotju]

/bot skip --comment "failed test already waived, shall be fixed later"

[tensorrt-cicd]

PR_Github #59699 [ skip ] triggered by Bot. Commit: a9d549c Link to invocation

[tensorrt-cicd]

PR_Github #59696 [ run ] completed with state ABORTED. Commit: a9d549c

Link to invocation

[tensorrt-cicd]

PR_Github #59699 [ skip ] completed with state SUCCESS. Commit: a9d549c
Skipping testing for commit a9d549c

Link to invocation