[None][test] reuse conftest mpi_pool_executor in MoE routing multi-GPU test

[xxi-nv]

Summary

Migrate the perfect-router multi-GPU routing test (test_perfect_router_load_balanced_multi_gpu) off a per-case MPIPoolExecutor onto the shared, module-scoped conftest mpi_pool_executor fixture (world_size=4 via indirect parametrization), the same pool other multi-GPU tests (test_moe_a2a / test_moe_comm) already use.

Details

• Before: each parametrized case did with MPIPoolExecutor(max_workers=4), re-spawning 4 worker processes (re-import + CUDA/NCCL/comm init) per case.

• After: the 4 workers are spawned once for the whole module and reused across all cases; a per-case _reset_perfect_router_comm_state() clears the perfect-router logits cache and NVLink one-sided workspace singletons between cases to match the fresh-process semantics of the previous per-case executor.

• The unused from mpi4py.futures import MPIPoolExecutor import is removed; @pytest.mark.threadleak(enabled=False) is added (the module-scoped pool manager thread persists by design).

Test plan

• pre-commit run --files tests/unittest/_torch/modules/test_moe_routing.py
• GB200 (sm100), origin/main matching container: pytest test_moe_routing.py::test_perfect_router_load_balanced_multi_gpu -> 72 passed in 381.26s.

Summary by CodeRabbit

• Tests
  • Improved reliability of the multi-GPU routing test by reusing a shared worker pool fixture instead of creating one directly.
  • Added cleanup between worker runs to prevent cross-test interference and ensure consistent results across repeated cases.
  • Updated the test to validate mapped worker executions explicitly.

[coderabbitai]

📝 Walkthrough

Walkthrough

The test module for MoE perfect-router now uses a shared, parametrized mpi_pool_executor fixture instead of creating a local MPIPoolExecutor. New helper functions reset comm/NVLink state and CUDA cache after each worker run to prevent cross-test contamination when worker processes are reused.

Changes

Perfect-Router Multi-GPU Test Refactor

Layer / File(s)	Summary
Worker comm state reset and entry wrapper tests/unittest/_torch/modules/test_moe_routing.py	Adds _reset_perfect_router_comm_state() to clear cached logits and NVLink one-sided workspace state and sync/empty CUDA cache; adds _perfect_router_worker_entry() that always resets state in a finally block after running the worker; updates test decorators to use mpi_pool_executor (parametrized indirect, 4 workers) and disables threadleak detection.
Test body dispatch via shared executor tests/unittest/_torch/modules/test_moe_routing.py	Replaces local MPIPoolExecutor creation with mpi_pool_executor.num_workers for world_size and dispatches work via mpi_pool_executor.map(_perfect_router_worker_entry, ...), asserting each result is None.

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

Possibly related PRs

• NVIDIA/TensorRT-LLM#15744: Both PRs modify MoE multi-GPU MPI-based tests to reuse pooled worker processes with per-task reset logic for cached one-sided comm/NVLink state.

Suggested reviewers: leslie-fang25

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is specific, concise, and accurately summarizes the main change to reuse the shared MPI pool in the MoE routing test.
Description check	✅ Passed	The description covers the change, rationale, and test plan, and is mostly complete despite missing the PR Checklist section.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

tests/unittest/_torch/modules/test_moe_routing.py (1)

891-944: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add explicit return annotations for the new functions.

The new helpers/test return None; annotate that and avoid returning the worker expression from _perfect_router_worker_entry. As per coding guidelines, “Always annotate functions. Make the return type None if the function does not return.”

Proposed annotation fix

-def _reset_perfect_router_comm_state():
+def _reset_perfect_router_comm_state() -> None:
@@
-def _perfect_router_worker_entry(*worker_args):
+def _perfect_router_worker_entry(*worker_args) -> None:
     """Run one perfect-router case, then reset state for the reused worker."""
     try:
-        return _perfect_router_worker(*worker_args)
+        _perfect_router_worker(*worker_args)
     finally:
         _reset_perfect_router_comm_state()
@@
 def test_perfect_router_load_balanced_multi_gpu(parallel_mode, routing_name,
                                                 num_tokens, dtype,
-                                                mpi_pool_executor):
+                                                mpi_pool_executor) -> None:

🤖 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/_torch/modules/test_moe_routing.py` around lines 891 - 944,
The new helpers in test_moe_routing should have explicit return type annotations
because they do not return a value. Update _reset_perfect_router_comm_state and
_perfect_router_worker_entry to be annotated as returning None, and change
_perfect_router_worker_entry so it does not return the result of
_perfect_router_worker but only runs it and then resets state in the finally
block. Keep the fix localized to these helper functions.

Source: Coding guidelines

🤖 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 `@tests/unittest/_torch/modules/test_moe_routing.py`:
- Around line 906-922: Synchronize the GPU before clearing shared routing state
in the cleanup path. In the test teardown around _perfect_router_worker, move
torch.cuda.synchronize() ahead of clearing _PERFECT_ROUTER_LOGITS_CACHE and
NVLinkOneSided workspace state, then keep the cache/workspace reset and
torch.cuda.empty_cache() afterward. Use the _PERFECT_ROUTER_LOGITS_CACHE and
NVLinkOneSided symbols to locate the cleanup block.

---

Nitpick comments:
In `@tests/unittest/_torch/modules/test_moe_routing.py`:
- Around line 891-944: The new helpers in test_moe_routing should have explicit
return type annotations because they do not return a value. Update
_reset_perfect_router_comm_state and _perfect_router_worker_entry to be
annotated as returning None, and change _perfect_router_worker_entry so it does
not return the result of _perfect_router_worker but only runs it and then resets
state in the finally block. Keep the fix localized to these helper functions.

🪄 Autofix (Beta)

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: fa2f8ad0-1e66-4075-ba4b-83acea0fd59e

📥 Commits

Reviewing files that changed from the base of the PR and between c81fa33 and a1ed65c.

📒 Files selected for processing (1)

• tests/unittest/_torch/modules/test_moe_routing.py

[xxi-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #57136 [ run ] triggered by Bot. Commit: 037a93e Link to invocation

[tensorrt-cicd]

PR_Github #57136 [ run ] completed with state SUCCESS. Commit: 037a93e
/LLM/main/L0_MergeRequest_PR pipeline #45919 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

[xxi-nv]

/bot run --disable-fail-fast

[xxi-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #57272 [ run ] triggered by Bot. Commit: 23c0f04 Link to invocation

[tensorrt-cicd]

PR_Github #57272 [ run ] completed with state SUCCESS. Commit: 23c0f04
/LLM/main/L0_MergeRequest_PR pipeline #46035 completed with status: 'SUCCESS'

CI Report

Link to invocation