From 299930bfc991719b34d79c1fe7368549ebd6151d Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Mon, 20 Jul 2026 09:14:27 +0000 Subject: [PATCH] [https://nvbugs/6480574][fix] mpi_session: guard MpiPoolSession.shutdown against partial init MpiPoolSession.shutdown() reads self.n_workers (in a log line) and self._wait_shutdown unguarded, but shutdown() is reachable on a session that never completed __init__: - from __del__ -> shutdown_abort() -> shutdown() on a partially-constructed object, and - from the released-session path, e.g. tests/unittest/executor/test_proxy_fast_death.py::test_pool_session_shutdown_never_blocks_after_release, which builds the object with MpiPoolSession.__new__(MpiPoolSession). In those states the attributes are absent and shutdown() raises AttributeError: 'MpiPoolSession' object has no attribute 'n_workers' (then '_wait_shutdown'). Guard both accesses with getattr(...), matching the existing _pool_dead guard, so a released or partially-constructed session shuts down cleanly instead of raising. No behavior change for fully-initialized sessions (n_workers / _wait_shutdown are always set by __init__; the guards only affect objects that bypass it). Also remove the temporary waiver for that test (nvbugs/6480574) so CI re-validates the now-fixed test. Signed-off-by: Yao Yao --- tensorrt_llm/llmapi/mpi_session.py | 11 +++++++---- tests/integration/test_lists/waives.txt | 1 - 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/llmapi/mpi_session.py b/tensorrt_llm/llmapi/mpi_session.py index 8f9264923782..b9d336a55150 100644 --- a/tensorrt_llm/llmapi/mpi_session.py +++ b/tensorrt_llm/llmapi/mpi_session.py @@ -61,6 +61,7 @@ def is_initialized() -> bool: def external_mpi_comm_available(model_world_size: int) -> bool: """Check if the current process is launched by mpirun and does not use MPIPoolExecutor to spawn processes. + e.g. mpirun -np 4 python script.py """ if ENABLE_MULTI_DEVICE: @@ -241,7 +242,9 @@ def __init__(self, n_workers: int, wait_shutdown: bool = False, env_overrides: Optional[Dict[str, str]] = None): - """Args: + """Create a pool session that spawns and manages MPI worker processes. + + Args: n_workers: number of MPI workers to spawn. wait_shutdown: when True, ``shutdown()`` blocks until the spawned worker processes have actually exited. ``MPIPoolExecutor.shutdown`` @@ -292,12 +295,12 @@ def shutdown(self, wait=True): wait = False if self.mpi_pool is not None: logger.info( - f"MpiPoolSession.shutdown: joining {self.n_workers} worker(s) " - f"(wait={wait})") + "MpiPoolSession.shutdown: joining " + f"{getattr(self, 'n_workers', '?')} worker(s) (wait={wait})") self.mpi_pool.shutdown(wait=wait) logger.info("MpiPoolSession.shutdown: done") self.mpi_pool = None - if self._wait_shutdown: + if getattr(self, '_wait_shutdown', False): self._wait_workers_exit() def _collect_worker_identities(self) -> Tuple: diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index ac2ded84fd37..7e92cea4df57 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -414,7 +414,6 @@ unittest/_torch/visual_gen/test_attention_integration.py::test_sage_attention_se unittest/_torch/visual_gen/test_attention_integration.py::test_sage_attention_self_attention[int8-2-1560] SKIP (https://nvbugs/6198760) unittest/auto_deploy/multigpu/custom_ops SKIP (https://nvbugs/6403920) unittest/disaggregated/test_kv_transfer.py::test_transfer_worker_v2[tp4_pp1_to_tp2_pp2] SKIP (https://nvbugs/6426834) -unittest/executor/test_proxy_fast_death.py::test_pool_session_shutdown_never_blocks_after_release SKIP (https://nvbugs/6480574) unittest/executor/test_rpc.py::TestRpcCorrectness::test_incremental_task_async SKIP (https://nvbugs/5741476) unittest/executor/test_rpc_proxy.py SKIP (https://nvbugs/5605741) unittest/executor/test_rpc_worker.py SKIP (https://nvbugs/5605741)