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

[chienchunhung]

Summary

The MPI launcher imports TensorRT-LLM again just to send its shutdown message, adding teardown latency to the remote-session tests reported in NVBUG 6913451. Use a standalone authenticated stop sender to avoid that extra import.

Scope

• Install the stop helper beside the launcher, preserving the pickle/HMAC wire format and clearing the inherited session key.
• Use main's existing launcher flow and task/stop/server status precedence. The helper bounds socket send/linger to five/one seconds; it does not impose an end-to-end shutdown deadline.
• Add helper, launcher-routing/status and production Queue coverage to CPU CI, plus shutdown documentation.

Verification

• Current head 6c62b795: packaging fixture repaired; local packaging/helper/launcher tests 28 passed, 1 Linux-only skip; hooks and DCO passed. Full CI rerun requested; GitHub authorization is queued and Jenkins launch is pending.
• Main CI #64570 at 1906eaf1 failed. ARM: 39/42 passed initially; flashinfer_workspace, flashinfer_temporary_cleanup, and client_sys_path hit the unchanged 60-second watchdog, then passed on the MPI-suite retry. No affected tests were waived or skipped.
• x86 affected tests were unrun because of a build infrastructure failure. No paired baseline comparison was completed; NVBUG resolution remains unverified.

Notes

Adapts the lightweight sender from #20293. #20232 separately ports lifecycle safeguards and removes main's eight-second polling delay; this PR leaves that polling behavior unchanged. Verification and maintainer review remain pending.

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

Dev Engineer Review

• The launcher now sends the authenticated MPI stop message through an adjacent standalone helper. The helper avoids importing TensorRT-LLM and MPI, consumes the inherited HMAC key, and preserves the pickle/HMAC message format.
• The helper limits ZeroMQ send wait to five seconds and socket linger to one second. These limits do not bound Python startup or server shutdown.
• The launcher retains its existing task, stop, and server status precedence. Packaging installs the helper beside the launcher.
• Review severity counts are unavailable from the supplied findings. Test results and overall completion status are not established by the changed files.

QA Engineer Review

• Added test_mpi_stop.py and test_mpi_stop_launcher.py; modified test_mpi_session.py and test_serve_schema_packaging.py.
• Added the stop-helper modules to l0_cpu.yml. test_mpi_session.py is listed in both l0_cpu.yml and l0_a100.yml with ISOLATION.
• Coverage includes authenticated stop delivery, invalid configuration, bounded send failure, key removal, launcher routing and status precedence, PAIR queue integration, and wheel script packaging.
• Coverage verdict: sufficient for the described helper and launcher behavior. The stated timeout limits do not cover end-to-end shutdown duration.

Per-File QA Perspective

• docs/source/llm-api/index.md: Documents remote MPI shutdown and timeout limits. Verify the troubleshooting guidance matches helper behavior.
• setup.py: Installs the stop helper alongside the launcher. Verify both scripts are available together in installed layouts.
• tensorrt_llm/llmapi/_llmapi_mpi_stop.py: Adds authenticated stop delivery without importing TensorRT-LLM or MPI. Verify key validation and clearing, message compatibility, and bounded ZeroMQ failure behavior.
• tensorrt_llm/llmapi/trtllm-llmapi-launch: Routes shutdown through the adjacent helper instead of the package module. Verify source and installed layouts preserve routing and status precedence.
• tests/integration/test_lists/test-db/l0_cpu.yml: Adds test_mpi_stop.py and test_mpi_stop_launcher.py; it also lists test_mpi_session.py with ISOLATION.
• tests/unittest/llmapi/test_mpi_session.py: Adds subprocess coverage for stop delivery to a PAIR server queue. The module is listed with ISOLATION in l0_cpu.yml and l0_a100.yml.
• tests/unittest/llmapi/test_mpi_stop.py: Covers authenticated delivery, import isolation, send timeout, configuration errors, and key clearing. The module is listed in l0_cpu.yml.
• tests/unittest/llmapi/test_mpi_stop_launcher.py: Covers launcher routing in source and installed layouts, environment handling, and status precedence. The module is listed in l0_cpu.yml.
• tests/unittest/scripts/test_serve_schema_packaging.py: Verifies that the wheel includes the launcher and stop helper scripts. Its test-list placement is not established by the supplied inspection results.

[chienchunhung]

/bot run

[tensorrt-cicd]

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

[tensorrt-cicd]

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]

/bot run --disable-fail-fast

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/TensorRT-LLM/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 735a94d5-6bad-4b4a-a33e-a1648c63872d

📥 Commits

Reviewing files that changed from the base of the PR and between 1906eaf and 6c62b79.

📒 Files selected for processing: 1

• 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 launcher now uses a standalone helper for authenticated MPI shutdown. The helper validates and clears the HMAC key, then sends a signed serialized None over ZeroMQ with bounded send and linger settings. Packaging, tests, and troubleshooting documentation cover the helper and launcher flow.

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, docs/source/llm-api/index.md	Adds the standalone sender, HMAC key cleanup, signed sentinel delivery, and bounded send settings. Tests exercise key handling, delivery, failures, and MPI session compatibility. Documentation describes dependencies and timeout limits.
Launcher routing and packaging tensorrt_llm/llmapi/trtllm-llmapi-launch, setup.py, tests/unittest/llmapi/test_mpi_stop_launcher.py, tests/unittest/scripts/test_serve_schema_packaging.py, tests/integration/test_lists/test-db/l0_cpu.yml	Routes shutdown through the adjacent helper and installs the helper script. Tests cover source and installed launcher behavior, environment handling, locking, status handling, and wheel contents.

Priority: ➖ Normal

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

Change: Bug fix

Sequence Diagram(s)

sequenceDiagram
  participant Launcher
  participant StopHelper
  participant MPICommunicationServer
  Launcher->>StopHelper: Execute adjacent shutdown script
  StopHelper->>MPICommunicationServer: Send HMAC-signed serialized None

Loading

Suggested reviewers: chzblych

Merge Risk: ⚪ Minimal · up to 6c62b

The wheel includes the shutdown helper alongside the launcher, so this packaging-test change presents no identified merge-blocking risk.

Pre-merge checks: 5

Check name	Status	Explanation
Docstring Coverage		Docstring coverage is 93.33% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 15 functions across 6 files.
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.
Title check		The title follows the required NVBugs and type format and clearly summarizes the main change: avoiding heavy imports during MPI shutdown.
Description check		The description explains the problem, solution, scope, testing, CI status, and checklist items. It uses equivalent headings instead of the template headings, but it provides the required information a…

Full details: Description check

Explanation

The description explains the problem, solution, scope, testing, CI status, and checklist items. It uses equivalent headings instead of the template headings, but it provides the required information and is mostly complete.

✨ 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]

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

[tensorrt-cicd]

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

[chienchunhung]

/bot run --disable-fail-fast

[tburt-nv]

Maybe we should replace trtllm-llmapi-launch with a python script?

[chienchunhung]

Maybe we should replace trtllm-llmapi-launch with a python script?

Agreed; a standalone Python launcher could simplify this considerably. I’d keep this PR focused on removing the heavy imports from shutdown and handle the launcher conversion separately, preserving its MPI environment isolation, workspace locking, and exit-status behavior. The Python launcher should also avoid importing TensorRT-LLM itself.

[tensorrt-cicd]

PR_Github #78527 [ run ] triggered by Bot. Commit: 6c62b79 Link to invocation

[tensorrt-cicd]

PR_Github #78527 [ run ] completed with state FAILURE. Commit: 6c62b79
/LLM/main/L0_MergeRequest_PR pipeline #64716 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]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #78581 [ run ] triggered by Bot. Commit: 6c62b79 Link to invocation

[tensorrt-cicd]

PR_Github #78581 [ run ] completed with state SUCCESS. Commit: 6c62b79
/LLM/main/L0_MergeRequest_PR pipeline #64736 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