[https://nvbugs/6435642][fix] handle session reuse worker registration (NVIDIA#16444)

Author: chienchunhung

tensorrt_llm/executor/proxy.py

@@ -573,7 +573,18 @@ def mpi_done_callback(future: concurrent.futures.Future):
             raise RuntimeError(
                 "Executor worker returned error") from ready_signal
 
-        if isinstance(self.mpi_session, MpiPoolSession) and len(status) == 3:
+        self._register_worker_processes(status)
+
+    def _register_worker_processes(self, status: tuple) -> None:
+        """Register identities returned by locally spawned MPI workers.
+
+        Test session reuse replaces this module's ``MpiPoolSession`` class
+        reference with a factory, so identify pool-backed sessions by excluding
+        the external communication session types.
+        """
+        if not isinstance(
+                self.mpi_session,
+            (MpiCommSession, RemoteMpiCommSessionClient)) and len(status) == 3:
             worker_process_identities: List[WorkerProcessIdentity] = status[2]
             self._worker_process_monitor.register(worker_process_identities)
 

tests/integration/test_lists/waives.txt

@@ -447,9 +447,7 @@ test_e2e.py::test_multi_nodes_eval[MiniMax-M2-tp16-mmlu] SKIP (https://nvbugs/63
 test_e2e.py::test_multi_nodes_eval[MiniMax-M3-tp16-mmlu] SKIP (https://nvbugs/6373561)
 test_e2e.py::test_ptp_quickstart_advanced_deepseek_r1_w4afp8_8gpus[DeepSeek-R1-W4AFP8-DeepSeek-R1/DeepSeek-R1-W4AFP8] SKIP (https://nvbugs/5836830)
 unittest/_torch/misc/test_share_tensor.py::TestShareTensor::test_share_tensor_different_dtypes SKIP (https://nvbugs/6418021)
-unittest/_torch/modeling -k "modeling_out_of_tree" SKIP (https://nvbugs/6426847)
 unittest/_torch/modeling -k "modeling_qwen" SKIP (https://nvbugs/6433376)
-unittest/_torch/modeling/test_modeling_out_of_tree.py::TestOutOfTree::test_llm_api[True] SKIP (https://nvbugs/6426847)
 unittest/_torch/modeling/test_modeling_qwen3_5_vl.py::test_qwen35_dense_vl_resolves_mamba_ssm_cache_dtype SKIP (https://nvbugs/6433376)
 unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend[act=Relu2-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=NVFP4-routing=Renormalize] SKIP (https://nvbugs/5989912)
 unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "TRTLLM" SKIP (https://nvbugs/6464169)

tests/unittest/executor/test_proxy_fast_death.py

@@ -16,10 +16,12 @@
 
 import asyncio
 import queue as _queue
+from unittest.mock import Mock
 
 import pytest
 
 from tensorrt_llm.executor import EngineDeadError
+from tensorrt_llm.executor import proxy as proxy_module
 from tensorrt_llm.executor.proxy import GenerationExecutorProxy
 from tensorrt_llm.executor.result import GenerationResult
 
@@ -97,6 +99,20 @@ def test_handle_worker_death_broadcasts_event_driven():
     assert proxy._error_queue.get_nowait() is cause
 
 
+def test_register_worker_processes_with_session_reuse_factory(monkeypatch):
+    """Session reuse replaces proxy.MpiPoolSession with a factory function."""
+    pool_session = object()
+    monkeypatch.setattr(proxy_module, "MpiPoolSession", lambda n_workers: pool_session)
+    proxy = _bare_proxy()
+    proxy.mpi_session = proxy_module.MpiPoolSession(1)
+    proxy._worker_process_monitor = Mock()
+    identities = [object()]
+
+    proxy._register_worker_processes((proxy.READY_SIGNAL, None, identities))
+
+    proxy._worker_process_monitor.register.assert_called_once_with(identities)
+
+
 def test_result_step_raises_on_engine_dead():
     res = GenerationResult.__new__(GenerationResult)
     res.queue = _queue.Queue()