From a9d549c1f4be02fbbd9090a8535dc1bf806e05d6 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Wed, 15 Jul 2026 10:06:06 -0700 Subject: [PATCH] [nvbugs/6435642][fix] handle session reuse worker registration Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/executor/proxy.py | 13 ++++++++++++- tests/integration/test_lists/waives.txt | 2 -- tests/unittest/executor/test_proxy_fast_death.py | 16 ++++++++++++++++ 3 files changed, 28 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index 11451bad5b67..4850bfcf2f5d 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/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) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 618015a23d2b..d44b5575af1b 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -433,9 +433,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/tests_lora_modules/test_nemotron_h_lora_sanity.py::TestNemotronHLoRA::test_lora_pp2_sanity SKIP (https://nvbugs/6428124) diff --git a/tests/unittest/executor/test_proxy_fast_death.py b/tests/unittest/executor/test_proxy_fast_death.py index ed0e62ae063a..d6bbac01544f 100644 --- a/tests/unittest/executor/test_proxy_fast_death.py +++ b/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()