diff --git a/tensorrt_llm/_torch/pyexecutor/hang_detector.py b/tensorrt_llm/_torch/pyexecutor/hang_detector.py index 2ae692ed5902..b4fbe2cb3899 100644 --- a/tensorrt_llm/_torch/pyexecutor/hang_detector.py +++ b/tensorrt_llm/_torch/pyexecutor/hang_detector.py @@ -17,6 +17,7 @@ import signal import sys import threading +import time from contextlib import contextmanager from typing import Callable, Optional @@ -26,6 +27,11 @@ # 137 == 128 + SIGKILL(9): the exit code a shell reports for a SIGKILL'd process. _HARD_KILL_EXIT_CODE = 137 +# Grace (seconds) between a rank's executor-loop crash and the hard kill of the +# whole world. Negative disables the kill entirely (escape hatch). +RANK_CRASH_KILL_GRACE_ENV = "TLLM_RANK_CRASH_HARD_KILL_GRACE" +_RANK_CRASH_KILL_GRACE_DEFAULT = 10.0 + def _best_effort_flush_streams() -> None: """Flush stdout/stderr without ever raising; diagnostics must not block hard kill.""" @@ -80,6 +86,100 @@ def propagate_hard_kill(exit_code: int = _HARD_KILL_EXIT_CODE) -> None: os.kill(os.getpid(), signal.SIGKILL) +def _rank_crash_kill_grace() -> Optional[float]: + """Resolve the crash-kill grace period; ``None`` means the kill is disabled.""" + raw = os.environ.get(RANK_CRASH_KILL_GRACE_ENV) + if raw is None: + return _RANK_CRASH_KILL_GRACE_DEFAULT + try: + grace = float(raw) + except ValueError: + _best_effort_log_error( + f"Invalid {RANK_CRASH_KILL_GRACE_ENV}={raw!r}; " + f"using default {_RANK_CRASH_KILL_GRACE_DEFAULT}s" + ) + return _RANK_CRASH_KILL_GRACE_DEFAULT + return None if grace < 0 else grace + + +def hard_kill_on_rank_crash(world_size: int) -> bool: + """Hard-kill the whole world after this rank's executor loop crashed. + + A rank whose executor loop died on an exception can never rejoin its + peers' collectives: without an explicit kill, every peer blocks in its + next collective until its own HangDetector fires (300 s), and the whole + test session burns that long for an error that was already known. + + The grace sleep before the kill is load-bearing: it gives the crashed + rank's cleaner error paths time to win the race, so the client reports + the ORIGINAL exception instead of a bare worker death — + - rank-local response waiters woken by the executor-loop cleanup read + the stashed error and surface it through the response path; + - during init, the worker's ready handshake returns the real error to + the proxy before the abort tears the world down; + - the worker main thread returning lets its mpi4py future complete with + the original exception. + + Never raises (it runs in a ``finally`` where an exception would mask the + original loop error). Returns True when the kill path was taken — only + observable in tests, where ``propagate_hard_kill`` is stubbed; in + production that call does not return. + """ + try: + if world_size <= 1: + # No peers to unblock; the worker's own death already completes + # its future/handshake with the original exception. + return False + grace = _rank_crash_kill_grace() + if grace is None: + return False + _best_effort_log_error( + f"Executor loop crashed on this rank; hard-killing all " + f"{world_size} ranks in {grace}s (peers cannot make progress " + f"without this rank). Set {RANK_CRASH_KILL_GRACE_ENV}=-1 to disable." + ) + if grace > 0: + time.sleep(grace) + propagate_hard_kill() + return True + except Exception as e: # noqa: BLE001 - must not mask the loop's original error + _best_effort_log_error(f"hard_kill_on_rank_crash failed (ignored): {e!r}") + return False + + +def start_rank_crash_kill_watchdog(world_size: int) -> Optional[threading.Thread]: + """Arm a daemon thread that hard-kills the world once the grace elapses. + + Must be armed BEFORE executor-loop cleanup: cleanup can block without + bound (e.g. ``wait()`` on a pending PP send handle wedged by the crash), + and a kill placed after it would never be reached — leaving peers to + burn in their own 300 s HangDetectors, the exact failure this kill + exists to avoid. The thread reuses ``hard_kill_on_rank_crash``, so the + kill fires at crash + grace whether cleanup finishes, blocks, or raises. + + Never raises. Returns the armed thread, or ``None`` when the kill is + not applicable (single rank, disabled by env) or the thread could not + be started — in that case the caller's post-cleanup kill remains the + only mechanism. + """ + try: + if world_size <= 1: + return None + if _rank_crash_kill_grace() is None: + return None + watchdog = threading.Thread( + target=hard_kill_on_rank_crash, + args=(world_size,), + name="rank_crash_kill_watchdog", + daemon=True, + ) + watchdog.start() + return watchdog + except Exception as e: # noqa: BLE001 - must not mask the loop's original error + _best_effort_log_error(f"failed to arm rank-crash kill watchdog (ignored): {e!r}") + return None + + class HangDetector: """Watchdog that fires when the executor loop stops checkpointing. diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 9253b2f41fbd..9992c85f84fb 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -66,7 +66,8 @@ from .guided_decoder import GuidedDecoder from .handle_additional_outputs import HandleAdditionalOutputs from .handle_logits import HandleLogits -from .hang_detector import HangDetector, propagate_hard_kill +from .hang_detector import (HangDetector, hard_kill_on_rank_crash, + propagate_hard_kill, start_rank_crash_kill_watchdog) from .kv_cache_manager_v2 import KVCacheManagerV2 from .kv_cache_stats import append_kv_cache_iteration_stats from .kv_cache_transceiver import (KvCacheTransceiver, @@ -1180,6 +1181,7 @@ def _flush_iter_stats_synced(self): # Performance metrics methods are in PerfMetricsManager (self.perf_manager) def _event_loop_wrapper(self): + crashed = False try: # Skip line profiler during warmup/memory estimation phase to avoid # saving incomplete results that would be overwritten anyway @@ -1189,6 +1191,7 @@ def _event_loop_wrapper(self): customized_gc_thresholds(self.garbage_collection_gen0_threshold): self.event_loop() except Exception as e: + crashed = True logger.error(f"Error in event loop: {e}") logger.error(traceback.format_exc()) # Stash the original error so local consumers @@ -1201,7 +1204,25 @@ def _event_loop_wrapper(self): self._event_loop_error = e raise e finally: - self._executor_loop_cleanup() + if crashed: + # Armed BEFORE cleanup: cleanup can block without bound on a + # send handle wedged by the crash, and a kill placed only + # after it would never fire. + start_rank_crash_kill_watchdog(self.dist.world_size) + try: + self._executor_loop_cleanup() + finally: + if crashed: + # Peers cannot make progress without this rank's loop: + # they would block in their next collective until their + # own HangDetectors fire 300s later. Kill the world now + # instead; the grace inside lets the stashed error reach + # rank-local waiters and the ready handshake first, so + # the client sees the original exception rather than a + # bare worker death. Nested finally: the kill must fire + # even when cleanup itself raises; the watchdog above + # covers cleanup blocking. + hard_kill_on_rank_crash(self.dist.world_size) @property def is_warmup(self) -> bool: diff --git a/tests/unittest/_torch/executor/test_hang_detector_kill.py b/tests/unittest/_torch/executor/test_hang_detector_kill.py index b4f122bb8ad3..709c0c3618e3 100644 --- a/tests/unittest/_torch/executor/test_hang_detector_kill.py +++ b/tests/unittest/_torch/executor/test_hang_detector_kill.py @@ -14,13 +14,24 @@ # limitations under the License. """HangDetector timer behavior and the hard-kill propagation mechanism (no GPU).""" +import contextlib import os import signal import subprocess import sys +import threading import time +import types -from tensorrt_llm._torch.pyexecutor.hang_detector import HangDetector +import pytest + +from tensorrt_llm._torch.pyexecutor import hang_detector as hd_module +from tensorrt_llm._torch.pyexecutor.hang_detector import ( + RANK_CRASH_KILL_GRACE_ENV, + HangDetector, + hard_kill_on_rank_crash, + start_rank_crash_kill_watchdog, +) def test_detector_fires_after_timeout(): @@ -82,3 +93,201 @@ def test_propagate_hard_kill_self_sigkills_without_mpi(): f"expected self-SIGKILL (-9), got {proc.returncode}; " f"stderr={proc.stderr.decode(errors='replace')[-500:]}" ) + + +# -------------------------------------------------------------------------- +# hard_kill_on_rank_crash: a rank whose executor loop crashed must kill the +# world (after a grace) instead of leaving peers to burn 300s in collectives. +# -------------------------------------------------------------------------- + + +def test_rank_crash_kill_single_rank_is_noop(monkeypatch): + """No peers to unblock: the worker's own death already carries the error.""" + kills = [] + monkeypatch.setattr(hd_module, "propagate_hard_kill", lambda: kills.append(1)) + assert hard_kill_on_rank_crash(world_size=1) is False + assert kills == [] + + +def test_rank_crash_kill_fires_for_multi_rank(monkeypatch): + kills = [] + monkeypatch.setattr(hd_module, "propagate_hard_kill", lambda: kills.append(1)) + monkeypatch.setenv(RANK_CRASH_KILL_GRACE_ENV, "0") + assert hard_kill_on_rank_crash(world_size=4) is True + assert kills == [1] + + +def test_rank_crash_kill_sleeps_grace_before_kill(monkeypatch): + """The grace must elapse BEFORE the kill so cleaner error paths win the race.""" + order = [] + monkeypatch.setattr(hd_module, "propagate_hard_kill", lambda: order.append("kill")) + monkeypatch.setattr(hd_module.time, "sleep", lambda s: order.append(("sleep", s))) + monkeypatch.setenv(RANK_CRASH_KILL_GRACE_ENV, "2.5") + assert hard_kill_on_rank_crash(world_size=2) is True + assert order == [("sleep", 2.5), "kill"] + + +def test_rank_crash_kill_disabled_by_negative_grace(monkeypatch): + kills = [] + monkeypatch.setattr(hd_module, "propagate_hard_kill", lambda: kills.append(1)) + monkeypatch.setenv(RANK_CRASH_KILL_GRACE_ENV, "-1") + assert hard_kill_on_rank_crash(world_size=8) is False + assert kills == [] + + +def test_rank_crash_kill_invalid_grace_uses_default(monkeypatch): + """A malformed env value must not disable the kill (fail-safe default).""" + order = [] + monkeypatch.setattr(hd_module, "propagate_hard_kill", lambda: order.append("kill")) + monkeypatch.setattr(hd_module.time, "sleep", lambda s: order.append(("sleep", s))) + monkeypatch.setenv(RANK_CRASH_KILL_GRACE_ENV, "bogus") + assert hard_kill_on_rank_crash(world_size=2) is True + assert order == [("sleep", 10.0), "kill"] + + +def test_rank_crash_kill_never_raises(monkeypatch): + """It runs in a `finally`: raising would mask the loop's original error.""" + + def boom(): + raise RuntimeError("abort machinery broken") + + monkeypatch.setattr(hd_module, "propagate_hard_kill", boom) + monkeypatch.setenv(RANK_CRASH_KILL_GRACE_ENV, "0") + assert hard_kill_on_rank_crash(world_size=2) is False + + +# -------------------------------------------------------------------------- +# start_rank_crash_kill_watchdog: the kill must fire even when executor-loop +# cleanup never returns (e.g. blocked on a PP send handle wedged by the +# crash), so it is armed in a daemon thread BEFORE cleanup starts. +# -------------------------------------------------------------------------- + + +def test_watchdog_kills_while_caller_blocks(monkeypatch): + """The kill fires from the watchdog thread with no help from the caller.""" + killed = threading.Event() + monkeypatch.setattr(hd_module, "propagate_hard_kill", killed.set) + monkeypatch.setenv(RANK_CRASH_KILL_GRACE_ENV, "0") + + watchdog = start_rank_crash_kill_watchdog(world_size=2) + + assert watchdog is not None + assert watchdog.daemon # must never block interpreter exit + # The caller does nothing further (it would be blocked in cleanup); + # the kill must fire regardless. + assert killed.wait(timeout=30.0) + watchdog.join(timeout=30.0) + + +def test_watchdog_not_armed_for_single_rank(monkeypatch): + kills = [] + monkeypatch.setattr(hd_module, "propagate_hard_kill", lambda: kills.append(1)) + monkeypatch.setenv(RANK_CRASH_KILL_GRACE_ENV, "0") + assert start_rank_crash_kill_watchdog(world_size=1) is None + assert kills == [] + + +def test_watchdog_not_armed_when_disabled(monkeypatch): + kills = [] + monkeypatch.setattr(hd_module, "propagate_hard_kill", lambda: kills.append(1)) + monkeypatch.setenv(RANK_CRASH_KILL_GRACE_ENV, "-1") + assert start_rank_crash_kill_watchdog(world_size=8) is None + assert kills == [] + + +# -------------------------------------------------------------------------- +# Wiring: PyExecutor._event_loop_wrapper must invoke the kill on the crash +# path only, and only after local cleanup has woken rank-local waiters. +# -------------------------------------------------------------------------- + + +def _bare_executor(pe, monkeypatch, world_size): + # Neutralize the profiling/GC context managers: they are irrelevant to the + # crash path and must not depend on env/GC state in a unit test. + monkeypatch.setattr(pe, "host_profiler_context", lambda enable: contextlib.nullcontext()) + monkeypatch.setattr(pe, "customized_gc_thresholds", lambda threshold: contextlib.nullcontext()) + ex = pe.PyExecutor.__new__(pe.PyExecutor) + ex.dist = types.SimpleNamespace(world_size=world_size) + ex.garbage_collection_gen0_threshold = None + return ex + + +def _stub_kill_paths(pe, monkeypatch, events): + monkeypatch.setattr( + pe, "hard_kill_on_rank_crash", lambda world_size: events.append(("kill", world_size)) + ) + monkeypatch.setattr( + pe, + "start_rank_crash_kill_watchdog", + lambda world_size: events.append(("watchdog", world_size)), + ) + + +def test_event_loop_wrapper_kills_world_on_crash(monkeypatch): + from tensorrt_llm._torch.pyexecutor import py_executor as pe + + events = [] + _stub_kill_paths(pe, monkeypatch, events) + ex = _bare_executor(pe, monkeypatch, world_size=4) + ex._executor_loop_cleanup = lambda: events.append("cleanup") + + def crash(): + raise ValueError("boom") + + ex.event_loop = crash + + with pytest.raises(ValueError, match="boom"): + ex._event_loop_wrapper() + + # The watchdog is armed BEFORE cleanup (cleanup can block forever); + # cleanup wakes rank-local waiters (who read the stashed error) BEFORE + # the direct kill tears the world down. + assert events == [("watchdog", 4), "cleanup", ("kill", 4)] + assert isinstance(ex._event_loop_error, ValueError) + + +def test_event_loop_wrapper_kills_world_when_cleanup_raises(monkeypatch): + """The kill must not be skippable by a cleanup failure. + + Cleanup runs precisely when the process is already unhealthy; if its + exception aborted the finally block before the kill, peers would burn + 300s in their HangDetectors — the worst case is exactly when the kill + matters most. + """ + from tensorrt_llm._torch.pyexecutor import py_executor as pe + + events = [] + _stub_kill_paths(pe, monkeypatch, events) + ex = _bare_executor(pe, monkeypatch, world_size=4) + + def broken_cleanup(): + events.append("cleanup") + raise RuntimeError("cleanup exploded") + + ex._executor_loop_cleanup = broken_cleanup + + def crash(): + raise ValueError("boom") + + ex.event_loop = crash + + with pytest.raises(RuntimeError, match="cleanup exploded"): + ex._event_loop_wrapper() + + assert events == [("watchdog", 4), "cleanup", ("kill", 4)] + # The original loop error stays reachable for rank-local consumers. + assert isinstance(ex._event_loop_error, ValueError) + + +def test_event_loop_wrapper_no_kill_on_clean_exit(monkeypatch): + from tensorrt_llm._torch.pyexecutor import py_executor as pe + + events = [] + _stub_kill_paths(pe, monkeypatch, events) + ex = _bare_executor(pe, monkeypatch, world_size=4) + ex._executor_loop_cleanup = lambda: events.append("cleanup") + ex.event_loop = lambda: None + + ex._event_loop_wrapper() + + assert events == ["cleanup"] diff --git a/tests/unittest/executor/test_proxy_fast_death.py b/tests/unittest/executor/test_proxy_fast_death.py index c72ed8c21b1a..f428a7862d90 100644 --- a/tests/unittest/executor/test_proxy_fast_death.py +++ b/tests/unittest/executor/test_proxy_fast_death.py @@ -522,6 +522,10 @@ def test_pool_session_shutdown_never_blocks_after_release(): from tensorrt_llm.llmapi.mpi_session import MpiPoolSession session = MpiPoolSession.__new__(MpiPoolSession) + # __new__ bypasses __init__ (which would spawn real MPI workers), so the + # attributes shutdown() reads must be provided here. + session.n_workers = 2 + session._wait_shutdown = False pool = _Mock() pool._pool.thread = None # no real manager thread to deregister session.mpi_pool = pool