[https://nvbugs/6913451][fix] Avoid heavy imports in MPI shutdown by chienchunhung · Pull Request #20434 · NVIDIA/TensorRT-LLM

chienchunhung · 2026-10-15T17:57:09Z

Summary

Avoid importing TensorRT-LLM again just to stop the MPI launcher, reducing shutdown overhead implicated in NVBUGs 6913451, 6964779 and 6964799. Preserve main's acknowledged shutdown protocol and lifecycle supervision.

Scope

The standalone stop helper uses pyzmq, scrubs the session key, authenticates replies without unpickling them, and retries on fresh PAIR connections.
Install it beside the launcher and resolve launcher symlinks. An unavailable helper uses a logged module fallback; a started helper's failure propagates. Existing supervision remains unchanged: the 120-second default stop budget bounds retries, and acknowledgment confirms receipt rather than completed worker shutdown.
Remove the blanket remote-MPI waiver for NVBUG6913451 and the submit/NVBUG6964779 and task_kwargs/NVBUG6964799 waivers. This restores eligibility for five remote-MPI variants in ARM/x86 CPU CI and, after merge, on main. Keep the separate submit_sync waiver for NVBUG6988813.

Verification

Latest CI: 66504 SUCCESS at pre-rebase head 697dd78196. It reused both CPU stages from 66405, where raw JUnit/logs confirm 154 affected cases passed and 1 skipped per architecture (ARM/x86), with no selected failures or retries. Coverage includes helper, launcher, supervisor, packaging, production queue/credential compatibility and five remote-MPI variants. submit_sync was skipped under NVBUG6988813. Run 66405's DeepSeek-V4 CUDA OOM and H100 infrastructure failures remain part of its history; the subsequent successful rerun is not fresh CPU execution.

Bug mapping below refers to test_mpi_session.py::test_remote_mpi_session[variant] in 66405 at 697dd78196, using the then-current 60-second shell watchdog:

Bug	ARM/x86 execution evidence	Limit
6913451	`submit`, `flashinfer_workspace`, `flashinfer_temporary_cleanup`, `task_kwargs`, `client_sys_path` passed; logs show successful task, helper, workers and server exits.	Five of six variants executed; no new `submit_sync` evidence or paired baseline timing.
6964779	`submit` passed on both architectures without its specific waiver.	Single-run evidence does not establish intermittent-bug closure.
6964799	`task_kwargs` passed on both architectures without its specific waiver.	Single-run evidence does not establish intermittent-bug closure.

Current head b501a43274 on main ab95e84be: applicable hooks, DCO and self-review passed. A local check using the production parser, waiver matcher and CPU stage-selection code confirms five eligible remote-MPI variants per architecture, with only submit_sync skipped for NVBUG6988813. Local macOS ARM64/Python 3.13 helper, launcher, supervisor and packaging tests passed 108 cases with 1 Linux-only skip at 77f1e77cd0; the current revision changes only waiver selection relative to that run; all tested code and test/configuration files are unchanged. Native MPI/full CI were not rerun after this rebase. Main's merged #20467 changes MPI task distribution and replaces the shell watchdog with a 120-second Python subprocess bound; those upstream runtime/test changes are preserved; this PR removes the blanket waiver to resume coverage. Earlier native results do not verify this new combination or prove all three intermittent bugs resolved. Repeated comparable ARM task-exit-to-launcher-exit timing remains outstanding.

Notes

Main already includes #20608 for acknowledged stop delivery and #20232 for lifecycle supervision; no separate parent stack remains. Adapted from #20293.

After merge, update all three NVBUGs with the merged revision, corresponding evidence and remaining gaps; assess closure separately for each bug.

Reviewed the repository PR checklist; tests and documentation included. No new public API, dependency or ownership change.

chienchunhung · 2026-10-15T22:42:44Z

/bot run

tensorrt-cicd · 2026-10-15T22:50:11Z

PR_Github #78241 [ run ] triggered by Bot. Commit: 53bc92a Link to invocation

tensorrt-cicd · 2026-10-16T01:44:06Z

PR_Github #78241 [ run ] completed with state FAILURE. Commit: 53bc92a
/LLM/release-1.3/L0_MergeRequest_PR pipeline #272 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

chienchunhung · 2026-10-16T03:35:10Z

/bot run --disable-fail-fast

coderabbitai · 2026-10-16T03:37:18Z

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Last reviewed commit: 333d35f
Latest commit, not reviewed: b501a43

Use the following commands to manage reviews:

@coderabbitai resume to resume automatic reviews.
@coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

▶️ Resume reviews
🔍 Trigger review

No actionable comments were generated in the recent review.

ℹ️ Recent review info

⚙️ Run configuration

Source mode: Auto
Repository: NVIDIA/TensorRT-LLM
Branch: main
File: .coderabbit.yaml
Commit: 491fe18
Layers: Repository YAML
Review profile: CHILL
Plan: Enterprise
Run ID: 875dd501-be2f-406c-b753-d43c269b2923

📥 Commits

Reviewing files that changed from the base of the PR and between 05a5433 and 333d35f.

📒 Files selected for processing: 8

docs/source/llm-api/index.md
setup.py
tensorrt_llm/llmapi/trtllm-llmapi-launch
tests/integration/test_lists/test-db/l0_cpu.yml
tests/unittest/llmapi/test_mpi_launcher_shutdown.py
tests/unittest/llmapi/test_mpi_session.py
tests/unittest/llmapi/test_mpi_stop_launcher.py
tests/unittest/scripts/test_serve_schema_packaging.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

Walkthrough

The change adds a standalone authenticated MPI shutdown helper and routes launcher shutdown through it when available. It also updates installation, tests, and documentation for helper behavior, fallback conditions, timeouts, and exit-status handling.

Changes

MPI shutdown flow

Layer / File(s)	Summary
Standalone shutdown sender `tensorrt_llm/llmapi/_llmapi_mpi_stop.py`, `tests/unittest/llmapi/test_mpi_stop.py`, `tests/unittest/llmapi/test_mpi_session.py`	The helper sends authenticated shutdown requests over fresh ZeroMQ PAIR connections and accepts valid authenticated acknowledgments. Tests cover retry and timeout behavior, key handling, malformed replies, and both standalone and module clients.
Launcher routing and packaging `tensorrt_llm/llmapi/trtllm-llmapi-launch`, `setup.py`, `tests/unittest/llmapi/test_mpi_stop_launcher.py`, `tests/unittest/llmapi/test_mpi_launcher_shutdown.py`, `tests/unittest/scripts/test_serve_schema_packaging.py`, `tests/integration/test_lists/test-db/l0_cpu.yml`, `docs/source/llm-api/index.md`	The launcher resolves its path and selects the sibling helper when it exists and is readable; otherwise, it uses the module command. Tests cover routing, fallback, exit statuses, and packaging. The documentation describes the protocol and timeout behavior.

Priority: Normal

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Bug fix

Sequence Diagram(s)

sequenceDiagram
  participant Launcher
  participant StopHelper
  participant ZeroMqQueue_PAIR_server
  Launcher->>StopHelper: Start guarded shutdown helper
  StopHelper->>ZeroMqQueue_PAIR_server: Send authenticated shutdown request
  ZeroMqQueue_PAIR_server->>StopHelper: Return acknowledgment
  StopHelper->>Launcher: Return helper status

Possibly related PRs

NVIDIA/TensorRT-LLM#19664: Changes the related remote MPI shutdown protocol and launcher handling that this PR’s standalone client uses.

Merge Risk: Minimal · up to 333d3

No confirmed issue blocks the new shutdown path from merging after normal checks.

Pre-merge checks: 4

Check name	Status	Explanation
Title check		The title follows the required ticket and type format and clearly describes the primary change: avoiding heavy imports during MPI shutdown.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Description check		The pull request description is complete and relevant. It explains the problem, solution, scope, fallback and supervision behavior, test coverage, CI evidence, known limitations, related changes, and …

Full details: Description check

Explanation

The pull request description is complete and relevant. It explains the problem, solution, scope, fallback and supervision behavior, test coverage, CI evidence, known limitations, related changes, and checklist status. It does not use the exact template headings, but it provides the required information in equivalent sections.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create a new PR

Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

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

tensorrt-cicd · 2026-10-16T03:41:24Z

PR_Github #78337 [ run ] triggered by Bot. Commit: 1906eaf Link to invocation

tensorrt-cicd · 2026-10-16T09:17:44Z

PR_Github #78337 [ run ] completed with state FAILURE. Commit: 1906eaf
/LLM/main/L0_MergeRequest_PR pipeline #64570 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

tensorrt-cicd · 2026-10-20T21:52:46Z

PR_Github #79402 [ run ] triggered by Bot. Commit: 05a5433 Link to invocation

tensorrt-cicd · 2026-10-20T21:56:45Z

PR_Github #79359 [ run ] completed with state ABORTED. Commit: e9ebff2

Link to invocation

chienchunhung · 2026-10-20T22:03:46Z

Heads-up from #20608, which fixes a hang in the same launcher stop path. The comm server's control socket is a ZMQ PAIR socket. It takes a new client only after it has processed the previous client's disconnect, and while it waits for the last task's worker teardown, a single stop request sent right after the task exits is dropped. The leader then never shuts down and the job hangs after the test has finished. That's 2 of 2 runs of the Kimi-K3 DWDP16 perf-sanity lane on 16 GB300 ranks.

#20608 keeps mgmn_leader_node --action stop, but makes it:

resend on a new connection until the server replies shutdown-ack, which the server now sends as soon as it reads the request;

exit 1 if no ack ever comes.

A send-once helper in the launcher would bring the hang back, but a lightweight sender is fine as long as it keeps that loop: a new PAIR connection per attempt, poll for the HMAC-signed "shutdown-ack" frame (skipping other messages), LINGER 0 on close, about 300 attempts of 1 s, then exit 1. Whichever PR lands second should carry it over. I'm happy to help port it.

Thanks for the heads up @chenfeiz0326 . Updated to depend on #20608 and preserve its acknowledgment/retry protocol. Added reconnect and protocol compatibility coverage.

tensorrt-cicd · 2026-10-21T03:56:52Z

PR_Github #79402 [ run ] completed with state SUCCESS. Commit: 05a5433
/LLM/main/L0_MergeRequest_PR pipeline #65395 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

chienchunhung · 2026-10-22T18:17:05Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-10-22T18:26:16Z

PR_Github #79942 [ run ] triggered by Bot. Commit: 333d35f Link to invocation

tensorrt-cicd · 2026-10-23T07:00:09Z

PR_Github #79942 [ run ] completed with state FAILURE. Commit: 333d35f
/LLM/main/L0_MergeRequest_PR pipeline #65806 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

chienchunhung · 2026-10-23T19:50:30Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-10-23T19:59:57Z

PR_Github #80250 [ run ] triggered by Bot. Commit: 237640b Link to invocation

tensorrt-cicd · 2026-10-24T03:16:13Z

PR_Github #80250 [ run ] completed with state FAILURE. Commit: 237640b
/LLM/main/L0_MergeRequest_PR pipeline #66044 completed with status: 'FAILURE'

CI Report

⚠️ Resource Maintenance — Test Coverage Reduced:

Resource maintenance in LLM/main/L0_Test-x86_64-Multi-GPU #2924: B200 multi-GPU stages are paused to relieve the CI queue backlog (109 running / 459 pending on 2026-10-22). Owner: @sharabeshj-nvda. The CI team decides when to lift it.
Affected stages:
DGX_B200-2_GPUs-PyTorch-1
DGX_B200-4_GPUs-PyTorch-1
DGX_B200-4_GPUs-PyTorch-2
DGX_B200-4_GPUs-PyTorch-3
DGX_B200-4_GPUs-PyTorch-Ray-1
DGX_B200-8_GPUs-PyTorch-1
DGX_B200-8_GPUs-PyTorch-2
DGX_B200-8_GPUs-PyTorch-3
DGX_B200-8_GPUs-PyTorch-4
DGX_B200-8_GPUs-PyTorch-Ray-1

⚠️ 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

chienchunhung · 2026-10-26T16:46:17Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-10-26T16:54:21Z

PR_Github #80654 [ run ] triggered by Bot. Commit: 697dd78 Link to invocation

tensorrt-cicd · 2026-10-27T00:17:29Z

PR_Github #80654 [ run ] completed with state SUCCESS. Commit: 697dd78
/LLM/main/L0_MergeRequest_PR pipeline #66405 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

chienchunhung · 2026-10-27T01:31:45Z

/bot run --disable-fail-fast

chienchunhung · 2026-10-27T02:37:10Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-10-27T02:43:13Z

PR_Github #80763 [ run ] triggered by Bot. Commit: 697dd78 Link to invocation

tensorrt-cicd · 2026-10-27T05:11:20Z

PR_Github #80763 [ run ] completed with state SUCCESS. Commit: 697dd78
/LLM/main/L0_MergeRequest_PR pipeline #66504 completed with status: 'SUCCESS'

CI Report

Link to invocation

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>

chienchunhung · 2026-10-27T17:09:23Z

/bot skip --comment "Full CI pipeline passed earlier before the rebase."

tensorrt-cicd · 2026-10-27T17:16:57Z

PR_Github #80935 [ skip ] triggered by Bot. Commit: b501a43 Link to invocation

tensorrt-cicd · 2026-10-27T17:33:08Z

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

Link to invocation

github-actions Bot assigned chienchunhung Oct 15, 2026

chienchunhung force-pushed the dev/release13-lightweight-mpi-stop branch from 2634a5f to 53bc92a Compare October 15, 2026 21:30

chienchunhung force-pushed the dev/release13-lightweight-mpi-stop branch from 53bc92a to aa8e105 Compare October 15, 2026 23:09

chienchunhung changed the title ~~[https://nvbugs/6913451][fix] Avoid heavy imports in release MPI shutdown~~ [https://nvbugs/6913451][fix] Avoid heavy imports in MPI shutdown Oct 15, 2026

chienchunhung changed the base branch from release/1.3 to main October 15, 2026 23:09

chienchunhung force-pushed the dev/release13-lightweight-mpi-stop branch from aa8e105 to 1906eaf Compare October 15, 2026 23:12

chienchunhung marked this pull request as ready for review October 16, 2026 03:33

chienchunhung requested review from a team as code owners October 16, 2026 03:33

chienchunhung requested review from BowenFu, EmmaQiaoCh, YihuiLu512, arysef, brnguyen2, mzweilz, nv-guomingz and zhaoyuanh-nvidia October 16, 2026 03:33

BowenFu approved these changes Oct 16, 2026

View reviewed changes

chenfeiz0326 mentioned this pull request Oct 16, 2026

[https://nvbugs/6843484][fix] Keep multi-node aggregated perf-sanity jobs in one NVLink domain #20418

Merged

1 task done

chienchunhung force-pushed the dev/release13-lightweight-mpi-stop branch from 1906eaf to 6c62b79 Compare October 16, 2026 16:10

chienchunhung force-pushed the dev/release13-lightweight-mpi-stop branch from 05a5433 to 333d35f Compare October 22, 2026 18:16

chienchunhung enabled auto-merge (squash) October 22, 2026 19:21

yunruis mentioned this pull request Oct 24, 2026

[None][feat] Add the PrimTS native sparse MLA FMHA for DeepSeek-V4 generation and enable it by default on SM100 #20209

Merged

1 task done

chienchunhung force-pushed the dev/release13-lightweight-mpi-stop branch from 237640b to 697dd78 Compare October 26, 2026 16:45

chienchunhung force-pushed the dev/release13-lightweight-mpi-stop branch from 697dd78 to dddb527 Compare October 27, 2026 16:35

[https://nvbugs/6913451][fix] Avoid heavy imports in MPI shutdown

b501a43

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>

chienchunhung force-pushed the dev/release13-lightweight-mpi-stop branch from dddb527 to b501a43 Compare October 27, 2026 16:52

chienchunhung mentioned this pull request Oct 27, 2026

[https://nvbugs/6913451][fix] Cut remote MPI session submit_sync wait and launcher stop latency #20293

Closed

2 tasks done

chienchunhung merged commit df95eda into NVIDIA:main Oct 27, 2026
11 checks passed

chienchunhung mentioned this pull request Oct 28, 2026

[https://nvbugs/6988813][test] Unwaive remote MPI submit_sync #21205

Merged

1 task done

Conversation

chienchunhung commented Oct 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Scope

Verification

Notes

Uh oh!

chienchunhung commented Oct 15, 2026

Uh oh!

tensorrt-cicd commented Oct 15, 2026

Uh oh!

tensorrt-cicd commented Oct 16, 2026

Uh oh!

chienchunhung commented Oct 16, 2026

Uh oh!

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

Uh oh!

Reviews paused

Walkthrough

Changes

Sequence Diagram(s)

Uh oh!

tensorrt-cicd commented Oct 16, 2026

Uh oh!

tensorrt-cicd commented Oct 16, 2026

Uh oh!

tensorrt-cicd commented Oct 20, 2026

Uh oh!

tensorrt-cicd commented Oct 20, 2026

Uh oh!

chienchunhung commented Oct 20, 2026

Uh oh!

tensorrt-cicd commented Oct 21, 2026

Uh oh!

chienchunhung commented Oct 22, 2026

Uh oh!

tensorrt-cicd commented Oct 22, 2026

Uh oh!

tensorrt-cicd commented Oct 23, 2026

Uh oh!

chienchunhung commented Oct 23, 2026

Uh oh!

tensorrt-cicd commented Oct 23, 2026

Uh oh!

tensorrt-cicd commented Oct 24, 2026

Uh oh!

chienchunhung commented Oct 26, 2026

Uh oh!

tensorrt-cicd commented Oct 26, 2026

Uh oh!

tensorrt-cicd commented Oct 27, 2026

Uh oh!

chienchunhung commented Oct 27, 2026

Uh oh!

chienchunhung commented Oct 27, 2026

Uh oh!

tensorrt-cicd commented Oct 27, 2026

Uh oh!

tensorrt-cicd commented Oct 27, 2026

Uh oh!

chienchunhung commented Oct 27, 2026

Uh oh!

tensorrt-cicd commented Oct 27, 2026

Uh oh!

tensorrt-cicd commented Oct 27, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

10 participants

chienchunhung commented Oct 15, 2026 •

edited

Loading

coderabbitai Bot commented Oct 16, 2026 •

edited

Loading