[TRTLLM-13409][fix] make proxy shutdown non-blocking when the engine is dead

tensorrt_llm/executor/proxy.py

@@ -311,6 +311,18 @@ def _mark_engine_dead(self, error: Optional[BaseException] = None) -> None:
                 result.queue.put(dead_error)
             except Exception:  # noqa: BLE001 - a full/closed queue must not stop the sweep
                 pass
+        # Release the session's exit joins here, not at teardown: interpreter
+        # exit joins non-daemon threads before any teardown code runs, so a
+        # wedged pool manager thread must be deregistered while user code is
+        # still alive. Non-destructive, hence safe for unowned sessions.
+        release = getattr(getattr(self, 'mpi_session', None),
+                          'release_exit_joins', None)
+        if release is not None:
+            try:
+                release()
+            except Exception as e:  # noqa: BLE001 - best-effort cleanup
+                logger.debug(
+                    f"MPI session exit-join release failed (ignored): {e!r}")
 
     def _handle_worker_death(self, error: BaseException) -> None:
         """Event-driven worker-death handler.
@@ -639,7 +651,15 @@ def shutdown(self):
 
         logger_debug('Proxy.shutdown...\n', "yellow")
 
+        # An abruptly-killed worker world (MPI_Abort, SIGKILL, OOM) never
+        # completes its mpi4py futures: give them one short collective grace
+        # instead of blocking on each, and skip the ones still pending.
+        if self._engine_dead:
+            concurrent.futures.wait(self.mpi_futures, timeout=5.0)
+
         for f in self.mpi_futures:
+            if self._engine_dead and not f.done():
+                continue
             try:
                 f.result()
             except:
@@ -656,7 +676,11 @@ def shutdown(self):
         if self.dispatch_result_thread is not None and self.dispatch_result_thread.is_alive(
         ):
             self.dispatch_result_thread.stop()
-            self.dispatch_result_thread.join()
+            # With the engine dead, the shutdown sentinel will never arrive
+            # and the dispatcher may be blocked in a ZMQ recv forever: bound
+            # the join and leak the daemon thread.
+            self.dispatch_result_thread.join(
+                timeout=5.0 if self._engine_dead else None)
 
         # step3: finish all remaining work
 
@@ -674,7 +698,11 @@ def shutdown(self):
 
         self.workers_started = False
         if self._owns_mpi_session:
-            self.mpi_session.shutdown()
+            if self._engine_dead:
+                # Anything joining a dead worker world blocks forever.
+                self.mpi_session.abandon()
+            else:
+                self.mpi_session.shutdown()
 
         # Process the errors in-case error during shutting down the threads
         self._handle_background_error()

tensorrt_llm/llmapi/mpi_session.py

@@ -146,6 +146,56 @@ def shutdown_abort(self, grace: float = 60, reason=None):
         fut.set_result(None)
         killer.join()
 
+    def release_exit_joins(self):
+        """Mark the worker world dead and release anything that would join it.
+
+        Non-destructive, so it may be called by a component that does not
+        own the session. Must not tear the session down -- only ensure that
+        nothing (interpreter exit, a later blocking ``shutdown()`` by the
+        owner) waits forever on the dead world. Default: no-op.
+        """
+
+    def abandon(self):
+        """Tear the session down without waiting on a dead worker world."""
+        self.release_exit_joins()
+        self.shutdown(wait=False)
+
+
+def _abandon_mpi_pool_threads(mpi_pool) -> None:
+    """Let interpreter exit proceed despite a wedged pool manager thread.
+
+    When the worker world dies abruptly, the ``MPIPoolExecutor`` manager
+    thread stays blocked in an MPI call forever, and process exit hangs on
+    it twice: mpi4py's exit hook joins every registered manager thread, and
+    CPython joins every non-daemon thread. Deregister the thread from both;
+    it is reaped with the process.
+
+    Best-effort: the touched names are private to mpi4py (``THREADS_QUEUES``
+    in ``_lib``/3.x and ``_core``/4.x) and CPython
+    (``threading._shutdown_locks``, 3.9-3.12). Where a name is absent, that
+    mechanism is left alone and exit may still block on it.
+    """
+    thread = getattr(getattr(mpi_pool, '_pool', None), 'thread', None)
+    if thread is None:
+        return
+    # mpi4py's own exit hook (joins all registered manager threads).
+    for mod_name in ('mpi4py.futures._lib', 'mpi4py.futures._core'):
+        mod = sys.modules.get(mod_name)
+        registry = getattr(mod, 'THREADS_QUEUES', None) if mod else None
+        if registry is not None:
+            try:
+                registry.pop(thread, None)
+            except Exception as e:  # noqa: BLE001 - best-effort cleanup
+                logger.debug(f"THREADS_QUEUES cleanup failed (ignored): {e!r}")
+    # CPython's non-daemon thread join at interpreter shutdown.
+    tstate_lock = getattr(thread, '_tstate_lock', None)
+    shutdown_locks = getattr(threading, '_shutdown_locks', None)
+    if tstate_lock is not None and shutdown_locks is not None:
+        try:
+            shutdown_locks.discard(tstate_lock)
+        except Exception as e:  # noqa: BLE001 - best-effort cleanup
+            logger.debug(f"_shutdown_locks cleanup failed (ignored): {e!r}")
+
 
 class MpiPoolSession(MpiSession):
 
@@ -174,10 +224,19 @@ def submit_sync(self, task: Callable[..., T], *args, **kwargs) -> List[T]:
         return [future.result() for future in futures]
 
     def shutdown(self, wait=True):
+        if getattr(self, '_pool_dead', False):
+            # A dead pool can never be joined; never block on it, no matter
+            # what the caller asked for.
+            wait = False
         if self.mpi_pool is not None:
             self.mpi_pool.shutdown(wait=wait)
             self.mpi_pool = None
 
+    def release_exit_joins(self):
+        if self.mpi_pool is not None:
+            _abandon_mpi_pool_threads(self.mpi_pool)
+        self._pool_dead = True
+
     def abort(self):
         self.get_comm().Abort(1)
 

tests/unittest/executor/test_proxy_fast_death.py

@@ -16,7 +16,10 @@
 
 import asyncio
 import queue as _queue
+import time as _time
+from concurrent.futures import Future as _Future
 from unittest.mock import Mock
+from unittest.mock import Mock as _Mock
 
 import pytest
 
@@ -338,3 +341,191 @@ def test_proxy_check_remote_worker_death_marks_engine_dead():
     proxy2 = _bare_proxy()
     proxy2.mpi_session = object()
     assert proxy2._check_remote_worker_death() is False
+
+
+# --- Non-blocking teardown on a dead engine ---
+# An abruptly-killed worker world never completes its mpi4py futures and
+# never sends the result-queue shutdown sentinel, so an unbounded shutdown()
+# blocks forever on f.result(), the dispatcher join, and the pool join.
+# These tests pin the bounded-teardown behavior.
+
+
+def _teardown_proxy(engine_dead):
+    proxy = _bare_proxy()
+    proxy.workers_started = True
+    proxy.doing_shutdown = True  # skip pre_shutdown(); teardown path only
+    proxy._engine_dead = engine_dead
+    proxy._fatal_error = RuntimeError("worker died") if engine_dead else None
+    proxy.dispatch_result_thread = None
+    proxy.rpc_client = None
+    proxy.request_queue = _Mock()
+    proxy.worker_init_status_queue = _Mock()
+    proxy.result_queue = _Mock()
+    proxy._resource_governor_queue = None
+    proxy._owns_mpi_session = True
+    proxy.mpi_session = _Mock()
+    proxy._handle_background_error = lambda *a, **k: None
+    return proxy
+
+
+def test_shutdown_does_not_block_on_dead_engine():
+    """With the engine dead, never-completing futures must not hang teardown."""
+    proxy = _teardown_proxy(engine_dead=True)
+    pending = _Future()  # never completes: abrupt worker death
+    done = _Future()
+    done.set_exception(RuntimeError("captured by mpi_done_callback already"))
+    proxy.mpi_futures = [pending, done]
+
+    dispatcher = _Mock()
+    dispatcher.is_alive.return_value = True
+    proxy.dispatch_result_thread = dispatcher
+
+    start = _time.monotonic()
+    proxy.shutdown()
+    elapsed = _time.monotonic() - start
+
+    # One collective 5 s grace, not an unbounded f.result() per future.
+    assert elapsed < 30
+    assert not pending.done()
+    # The dispatcher join is bounded (daemon thread is leaked, not awaited).
+    dispatcher.stop.assert_called_once()
+    dispatcher.join.assert_called_once_with(timeout=5.0)
+    # The dead pool is not joined; the session is abandoned instead.
+    proxy.mpi_session.abandon.assert_called_once_with()
+    proxy.mpi_session.shutdown.assert_not_called()
+    assert proxy.workers_started is False
+
+
+def test_shutdown_keeps_blocking_semantics_when_engine_alive():
+    """Orderly shutdown is unchanged: futures reaped, session joined."""
+    proxy = _teardown_proxy(engine_dead=False)
+    done = _Future()
+    done.set_result(None)
+    proxy.mpi_futures = [done]
+
+    dispatcher = _Mock()
+    dispatcher.is_alive.return_value = True
+    proxy.dispatch_result_thread = dispatcher
+
+    proxy.shutdown()
+
+    dispatcher.join.assert_called_once_with(timeout=None)
+    proxy.mpi_session.shutdown.assert_called_once_with()
+    proxy.mpi_session.abandon.assert_not_called()
+    assert proxy.workers_started is False
+
+
+def test_shutdown_does_not_shut_down_external_session():
+    """An externally owned session must stay alive even on a dead engine."""
+    proxy = _teardown_proxy(engine_dead=True)
+    proxy._owns_mpi_session = False
+    proxy.mpi_futures = []
+
+    proxy.shutdown()
+
+    proxy.mpi_session.shutdown.assert_not_called()
+    proxy.mpi_session.abandon.assert_not_called()
+
+
+def test_abandon_mpi_pool_threads_unblocks_interpreter_exit():
+    """Both exit-join mechanisms release a wedged pool manager thread.
+
+    The thread is deregistered from mpi4py's THREADS_QUEUES and CPython's
+    _shutdown_locks, so process exit can proceed without joining it.
+    """
+    import sys as _sys
+    import threading as _threading
+    import types as _types
+
+    from tensorrt_llm.llmapi.mpi_session import _abandon_mpi_pool_threads
+
+    release = _threading.Event()
+    wedged = _threading.Thread(target=release.wait, name="fake_manager")
+    wedged.daemon = False
+    wedged.start()
+    try:
+        # Fake mpi4py registry module, as mpi4py would have registered it.
+        fake_mod = _types.ModuleType("mpi4py.futures._lib")
+        fake_mod.THREADS_QUEUES = {wedged: object()}
+        prev = _sys.modules.get("mpi4py.futures._lib")
+        _sys.modules["mpi4py.futures._lib"] = fake_mod
+        try:
+            fake_pool = _Mock()
+            fake_pool._pool.thread = wedged
+
+            _abandon_mpi_pool_threads(fake_pool)
+
+            assert wedged not in fake_mod.THREADS_QUEUES
+            shutdown_locks = getattr(_threading, "_shutdown_locks", None)
+            if shutdown_locks is not None:  # CPython 3.9-3.12
+                assert wedged._tstate_lock not in shutdown_locks
+        finally:
+            if prev is None:
+                del _sys.modules["mpi4py.futures._lib"]
+            else:
+                _sys.modules["mpi4py.futures._lib"] = prev
+    finally:
+        release.set()
+        wedged.join(timeout=5)
+
+
+def test_abandon_mpi_pool_threads_tolerates_missing_pool():
+    from tensorrt_llm.llmapi.mpi_session import _abandon_mpi_pool_threads
+
+    _abandon_mpi_pool_threads(None)
+    _abandon_mpi_pool_threads(object())  # no _pool attribute
+
+
+def test_mark_engine_dead_releases_exit_joins_immediately():
+    """Exit-join release must happen at detection time, not at teardown.
+
+    CPython's exit sequence joins non-daemon threads before atexit/GC can
+    run shutdown(), so a wedged pool manager thread must be deregistered
+    the moment the engine is marked dead.
+    """
+    proxy = _bare_proxy()
+    proxy.mpi_session = _Mock()
+
+    proxy._mark_engine_dead(RuntimeError("worker died"))
+    proxy.mpi_session.release_exit_joins.assert_called_once_with()
+
+    # Sticky: a second death report must not re-release.
+    proxy._mark_engine_dead(RuntimeError("again"))
+    proxy.mpi_session.release_exit_joins.assert_called_once_with()
+
+
+def test_mark_engine_dead_releases_external_sessions_too():
+    """Ownership must not gate the exit-join release.
+
+    The LLM API creates the session and passes it in, so the proxy does
+    not own it; the release is non-destructive bookkeeping and must still
+    happen.
+    """
+    proxy = _bare_proxy()
+    proxy._owns_mpi_session = False
+    proxy.mpi_session = _Mock()
+
+    proxy._mark_engine_dead(RuntimeError("worker died"))
+    proxy.mpi_session.release_exit_joins.assert_called_once_with()
+    # But destructive teardown is still reserved for the owner.
+    proxy.mpi_session.abandon.assert_not_called()
+    proxy.mpi_session.shutdown.assert_not_called()
+
+
+def test_pool_session_shutdown_never_blocks_after_release():
+    """A released session never blocks, even on an explicit shutdown().
+
+    After release_exit_joins(), a blocking shutdown() from the session
+    owner must not join the dead pool.
+    """
+    from tensorrt_llm.llmapi.mpi_session import MpiPoolSession
+
+    session = MpiPoolSession.__new__(MpiPoolSession)
+    pool = _Mock()
+    pool._pool.thread = None  # no real manager thread to deregister
+    session.mpi_pool = pool
+
+    session.release_exit_joins()
+    session.shutdown()  # owner asks for the default blocking shutdown
+
+    pool.shutdown.assert_called_once_with(wait=False)