diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index 4850bfcf2f5d..07d0ea0414fa 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/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() diff --git a/tensorrt_llm/llmapi/mpi_session.py b/tensorrt_llm/llmapi/mpi_session.py index 0fa5271f8acf..bbe82c09fe4c 100644 --- a/tensorrt_llm/llmapi/mpi_session.py +++ b/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) diff --git a/tests/unittest/executor/test_proxy_fast_death.py b/tests/unittest/executor/test_proxy_fast_death.py index d6bbac01544f..df9756307917 100644 --- a/tests/unittest/executor/test_proxy_fast_death.py +++ b/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)