diff --git a/tensorrt_llm/_torch/pyexecutor/hang_detector.py b/tensorrt_llm/_torch/pyexecutor/hang_detector.py index f6dde7c58959..2620d4aa4e13 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.""" @@ -44,6 +50,14 @@ def _best_effort_log_error(message: str) -> None: pass +def _best_effort_log_debug(message: str) -> None: + """Log at debug level without ever raising; diagnostics must not block hard kill.""" + try: + logger.debug(message) + except Exception: # noqa: BLE001 - diagnostics must not block hard kill + pass + + def propagate_hard_kill(exit_code: int = _HARD_KILL_EXIT_CODE) -> None: """Hard-kill this rank and propagate the kill to peer ranks. @@ -80,6 +94,224 @@ 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 _remaining_kill_grace(grace: float, deadline: Optional[float]) -> float: + """Time left before the kill must fire. + + ``deadline`` (a ``time.monotonic()`` stamp) lets a kill that was already + armed elsewhere keep its ORIGINAL fire time when it is handed over to + another waiter, so the handover cannot push the kill out by a second + grace. + + The ``max(0.0, ...)`` is belt-and-braces only: it keeps the return value + meaningful as "time left" for callers and logs. It is NOT what stops a + negative sleep -- ``_wait_out_kill_grace`` does that with its + ``remaining > 0`` guard (and ``Event.wait`` returns immediately for a + negative timeout anyway). Do not drop that guard on the strength of this + clamp. + """ + if deadline is None: + return grace + return max(0.0, deadline - time.monotonic()) + + +def _wait_out_kill_grace(remaining: float, cancelled: Optional[threading.Event]) -> bool: + """Sleep out the crash-kill grace; return False if the kill was cancelled. + + ``cancelled`` makes the wait interruptible so the timer can be handed + over to another waiter instead of two clocks running at once. + + The ``remaining > 0`` guard is load-bearing: a deadline already in the + past must fire the kill now, and ``time.sleep`` of a negative duration + would raise into ``hard_kill_on_rank_crash``'s blanket except and drop + the kill -- precisely in the case (cleanup outlasted the grace) the + watchdog exists for. + """ + if cancelled is None: + if remaining > 0: + time.sleep(remaining) + return True + return not cancelled.wait(remaining) + + +def hard_kill_on_rank_crash( + world_size: int, + deadline: Optional[float] = None, + cancelled: Optional[threading.Event] = None, + error_delivered: Optional[threading.Event] = None, +) -> 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. Returns False when the kill does + not apply (single rank, disabled by env) or was cancelled during the + grace. + """ + 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 + remaining = _remaining_kill_grace(grace, deadline) + _best_effort_log_error( + f"Executor loop crashed on this rank; hard-killing all " + f"{world_size} ranks in {remaining:g}s (peers cannot make progress " + f"without this rank). Set {RANK_CRASH_KILL_GRACE_ENV}=-1 to disable." + ) + if not _wait_out_kill_grace(remaining, cancelled): + # Debug, not error: the only caller that cancels does so to take + # the same kill over on the same deadline. Logging "cancelled" at + # ERROR right before the world is SIGKILLed reads during triage + # as "the kill was called off", which is the opposite of what + # happens. + _best_effort_log_debug( + "Rank-crash hard kill timer disarmed (handed over or no " + "longer needed); this timer will not fire." + ) + return False + # The grace has elapsed. Before killing, check whether the crash + # already reached the client. + # + # `crashed` upstream means "the loop raised before its break", which + # is broader than "peers are stranded". In a SYMMETRIC crash -- a + # deterministic Python error, a bad config, an OOM at the same batch -- + # every rank raises, nobody is stranded, and every rank arms this kill. + # Firing then would replace N clean tracebacks with a bare exit 137. + # + # The kill exists to stop peers blocking in a collective forever. If + # the stashed error has already been surfaced to the client, the + # failure is diagnosable and the kill buys nothing, so skip it and let + # the process exit normally with its original exception. + if error_delivered is not None and error_delivered.is_set(): + _best_effort_log_error( + "Rank-crash hard kill NOT fired: the executor-loop error " + "already reached the client, so the failure is reportable " + "without killing the world. Peers that are genuinely stranded " + "are still covered -- in that case nothing consumes the error " + "and this kill fires as before." + ) + return False + 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 + + +class RankCrashKillWatchdog(threading.Thread): + """Daemon thread that hard-kills the world once the crash grace elapses. + + A plain ``Thread`` for backwards compatibility (callers may still join it + or inspect ``daemon``), plus the two things needed to hand the timer over + to another waiter instead of running two clocks: + + - ``cancel()`` disarms THIS timer. It is a bookkeeping aid, not a safety + net: the only caller cancels in order to take the same kill over on the + same deadline one line later, so cancelling does not spare a rank. What + decides whether a rank is killed at all is the ``crashed`` predicate in + ``PyExecutor._event_loop_wrapper``. + - ``deadline`` exposes the original fire time so the caller that takes + over still fires at crash + grace rather than restarting the clock. + """ + + def __init__( + self, world_size: int, grace: float, error_delivered: Optional[threading.Event] = None + ): + super().__init__(name="rank_crash_kill_watchdog", daemon=True) + self._world_size = world_size + self.deadline = time.monotonic() + max(0.0, grace) + self._cancelled = threading.Event() + self._error_delivered = error_delivered + + def cancel(self) -> None: + """Disarm the kill. Never raises; safe to call more than once.""" + self._cancelled.set() + + @property + def cancelled(self) -> bool: + return self._cancelled.is_set() + + def run(self) -> None: + hard_kill_on_rank_crash( + self._world_size, + deadline=self.deadline, + cancelled=self._cancelled, + error_delivered=self._error_delivered, + ) + + +def start_rank_crash_kill_watchdog( + world_size: int, + error_delivered: Optional[threading.Event] = None, +) -> Optional[RankCrashKillWatchdog]: + """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. + + The returned watchdog is the ONLY timer while cleanup runs; the caller + is expected to ``cancel()`` it once cleanup returns and carry the kill + (with the same ``deadline``) itself, so the two paths never race with + two independent clocks. + + Never raises. Returns the armed watchdog, 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 + grace = _rank_crash_kill_grace() + if grace is None: + return None + watchdog = RankCrashKillWatchdog(world_size, grace, error_delivered) + 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 2fa5a5d0607b..7c45b29b24b3 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, @@ -637,6 +638,11 @@ def __init__( # broadcast an ErrorResponse to every pending request, waking # callers parked in queue.get() / aqueue.get(). self._event_loop_error: Optional[BaseException] = None + # Set once the stashed error has been surfaced to a client. Gates the + # rank-crash hard kill: if the failure is already reportable, killing + # the world only replaces a traceback with exit 137. threading.Event + # because the kill runs on a daemon thread. + self._event_loop_error_delivered = threading.Event() # kv cache events self.kv_cache_manager = self.resource_manager.resource_managers.get( @@ -834,6 +840,15 @@ def __init__( self.kv_cache_manager.snapshot_warmup_baseline() self.is_shutdown = False + # Set at the executor loops' normal-exit `break` sites, and ONLY + # there. It answers exactly one question for _event_loop_wrapper: + # "did event_loop() reach its own termination?" -- which decides + # whether an escaping exception stranded this rank's peers. + # is_shutdown cannot answer it: _handle_errors sets is_shutdown + # rank-locally on a fatal error (e.g. a CUDA illegal address on this + # rank alone) while peers are told nothing and the loop keeps running + # collectives, so a crash after that point still strands them. + self._event_loop_completed = False self._fatal_error: Optional[BaseException] = None self._error_budget = ErrorBudget() self._disagg_timed_out_ctx_cancelled_ids: set[int] = set() @@ -1190,6 +1205,8 @@ def _flush_iter_stats_synced(self): # Performance metrics methods are in PerfMetricsManager (self.perf_manager) def _event_loop_wrapper(self): + crashed = False + self._event_loop_completed = False try: # Skip line profiler during warmup/memory estimation phase to avoid # saving incomplete results that would be overwritten anyway @@ -1199,6 +1216,20 @@ def _event_loop_wrapper(self): customized_gc_thresholds(self.garbage_collection_gen0_threshold): self.event_loop() except Exception as e: + # A raise AFTER the loop reached its own normal-exit `break` (from + # the profiler's or hang detector's __exit__, or from the enclosing + # context managers) is a teardown error: this rank finished its + # work and no peer is waiting on it, so log it but never escalate + # to SIGKILLing the job. Anything else -- including a raise before + # the loop ever started -- leaves peers blocked in their next + # collective, which is what the kill exists to cut short. + # + # Deliberately NOT is_shutdown: _handle_errors flips that flag + # rank-locally on a fatal error (a CUDA illegal address on one rank + # is classified immediate_fatal and bypasses the error budget) and + # tells peers nothing, so a crash after that point -- the single + # most common trigger for this kill -- still strands them. + crashed = not self._event_loop_completed logger.error(f"Error in event loop: {e}") logger.error(traceback.format_exc()) # Stash the original error so local consumers @@ -1210,8 +1241,51 @@ def _event_loop_wrapper(self): # _executor_loop_cleanup is enough to wake local waiters. self._event_loop_error = e raise e + except BaseException: + # SystemExit / KeyboardInterrupt are NOT Exception, so they reach + # here rather than the handler above. Peers are just as stranded, + # but these are deliberate teardown signals -- the launcher is + # already tearing the job down -- and arming an MPI_Abort on top + # would turn a clean Ctrl-C into exit 137. Left unarmed on + # purpose; `crashed` stays False. Stated explicitly because the + # comment above describes the invariant the *Exception* path + # enforces, not this one. + raise finally: - self._executor_loop_cleanup() + # 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. + watchdog = start_rank_crash_kill_watchdog( + self.dist.world_size, + error_delivered=self._event_loop_error_delivered, + ) if crashed else None + 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. + # + # Cleanup returned, so the watchdog's only job (covering + # a cleanup that never returns) is done: hand the timer + # over rather than leave two running. This is bookkeeping, + # not protection -- the kill still fires, on the SAME + # deadline, one line below. Whether a rank is killed at + # all is decided solely by `crashed` above. + deadline = None + if watchdog is not None: + watchdog.cancel() + deadline = watchdog.deadline + hard_kill_on_rank_crash( + self.dist.world_size, + deadline=deadline, + error_delivered=self._event_loop_error_delivered) @property def is_warmup(self) -> bool: @@ -2551,6 +2625,7 @@ def _executor_loop_pp(self): # Fetch new requests from request queue new_requests = self._fetch_and_activate_new_requests() if self.should_stop_processing: + self._event_loop_completed = True break self._handle_control_request() @@ -2852,6 +2927,12 @@ def handle_executed_batches(executed_batch_num: int): self.iter_counter += 1 # Stage 5: Handle remaining executed batches in the queue. + # Note: _event_loop_completed was already set at the break above, + # so a raise in this drain is classified as benign teardown rather + # than a peer-stranding crash. That is deliberate: reaching here + # means every rank observed should_stop_processing, so no peer is + # parked in a collective waiting on this one -- this drain only + # consumes from a rank-local queue. while self.unhandled_batch_counter > 0: with nvtx_range("get_executed_batch"): executed_batch = self.executed_batch_response_queue.get() @@ -4033,6 +4114,7 @@ def _executor_loop(self): scheduled_batch, iter_stats = self._prepare_and_schedule_batch() if scheduled_batch is None: + self._event_loop_completed = True break can_forward, should_retry = self._check_benchmark_disagg_gate( @@ -4509,6 +4591,7 @@ def _executor_loop_overlap(self): scheduled_batch, iter_stats = self._prepare_and_schedule_batch() if scheduled_batch is None: + self._event_loop_completed = True break can_forward, should_retry = self._check_benchmark_disagg_gate( @@ -7081,6 +7164,9 @@ def key_has_response(): # instead of hanging here or hitting a KeyError below. error = self._event_loop_error if error is not None: + # The caller is about to see the original exception, so + # the crash is reportable without killing the world. + self._event_loop_error_delivered.set() raise RuntimeError( f"Event loop terminated with error: {error}") from error raise RuntimeError( diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 5e1413a4f0da..fabf1a657cbc 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -1152,6 +1152,15 @@ def __call__(self, timeout: Optional[float] = None) -> bool: # thread in that case too — see nvbug 6038228. error = getattr(self.worker.engine, "_event_loop_error", None) if error is not None: + # Broadcasting wakes every pending GenerationResult with the real + # error, so the crash is reportable. Tell the rank-crash kill that, + # so a symmetric crash (every rank raised the same deterministic + # error, nobody stranded) ends in N tracebacks rather than in + # MPI_Abort replacing them with a bare exit 137. + delivered = getattr(self.worker.engine, + "_event_loop_error_delivered", None) + if delivered is not None: + delivered.set() return self._broadcast_event_loop_error(error) return True diff --git a/tests/unittest/_torch/executor/test_hang_detector_kill.py b/tests/unittest/_torch/executor/test_hang_detector_kill.py index 0962df441cc6..541093f9a7d6 100644 --- a/tests/unittest/_torch/executor/test_hang_detector_kill.py +++ b/tests/unittest/_torch/executor/test_hang_detector_kill.py @@ -15,14 +15,25 @@ """HangDetector timer behavior and the hard-kill propagation mechanism (no GPU).""" import asyncio +import contextlib import os +import shutil import signal import subprocess import sys +import threading import time +import types + +import pytest from tensorrt_llm._torch.pyexecutor import hang_detector as hang_detector_module -from tensorrt_llm._torch.pyexecutor.hang_detector import HangDetector +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(): @@ -117,3 +128,887 @@ 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(hang_detector_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(hang_detector_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 _assert_slept_then_killed(order, grace): + """Assert the grace was slept out EXACTLY ONCE before EXACTLY ONE kill. + + Patching time.sleep is process-wide, so an unrelated background thread can + append its own ("sleep", x) while this runs; asserting exact list equality + would flake on that. But the counts must still be pinned: sleeping the + grace twice before killing (i.e. crash + 2*grace) is precisely the bug + class this PR series introduced with its two independent timers, and a + membership-only check accepts it. + """ + assert order.count(("sleep", grace)) == 1, order + assert order.count("kill") == 1, order + assert order.index(("sleep", grace)) < order.index("kill"), order + + +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(hang_detector_module, "propagate_hard_kill", lambda: order.append("kill")) + monkeypatch.setattr(hang_detector_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_slept_then_killed(order, 2.5) + + +def test_rank_crash_kill_disabled_by_negative_grace(monkeypatch): + kills = [] + monkeypatch.setattr(hang_detector_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(hang_detector_module, "propagate_hard_kill", lambda: order.append("kill")) + monkeypatch.setattr(hang_detector_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_slept_then_killed(order, 10.0) + + +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(hang_detector_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(hang_detector_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(hang_detector_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(hang_detector_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 == [] + + +def test_watchdog_cancel_disarms_this_timer(monkeypatch): + """cancel() must break the grace wait immediately, not after it elapses. + + This is the handover primitive, NOT protection against a spurious kill: + the only production caller cancels in order to take the same kill over on + the same deadline. What decides whether a rank is killed at all is the + `crashed` predicate in _event_loop_wrapper. + """ + kills = [] + monkeypatch.setattr(hang_detector_module, "propagate_hard_kill", lambda: kills.append(1)) + monkeypatch.setenv(RANK_CRASH_KILL_GRACE_ENV, "30") + + watchdog = start_rank_crash_kill_watchdog(world_size=2) + assert watchdog is not None + try: + watchdog.cancel() + watchdog.join(timeout=10.0) + assert not watchdog.is_alive() + assert kills == [] + assert watchdog.cancelled is True + finally: + # Never let an armed killer thread outlive the stubbed + # propagate_hard_kill: if cancel() ever regresses, the real one would + # SIGKILL the pytest process once monkeypatch restores it. + watchdog.cancel() + watchdog.join(timeout=60.0) + + +def test_watchdog_deadline_is_grace_from_arming(monkeypatch): + monkeypatch.setenv(RANK_CRASH_KILL_GRACE_ENV, "5") + watchdog = hang_detector_module.RankCrashKillWatchdog(world_size=2, grace=5.0) + assert watchdog.deadline == pytest.approx(time.monotonic() + 5.0, abs=0.5) + + +def test_kill_keeps_original_deadline_on_handover(monkeypatch): + """Handing the kill over must not restart the grace clock. + + The caller cancels the watchdog once cleanup returns and carries the kill + itself; passing the watchdog's deadline must make it sleep the REMAINING + time, not a fresh grace. Asserted on the exact duration handed to sleep + (with monotonic pinned) rather than on wall-clock, so the margin does not + depend on scheduling luck on a loaded CI node. + """ + order = [] + monkeypatch.setattr(hang_detector_module, "propagate_hard_kill", lambda: order.append("kill")) + monkeypatch.setattr(hang_detector_module.time, "sleep", lambda s: order.append(("sleep", s))) + monkeypatch.setattr(hang_detector_module.time, "monotonic", lambda: 1000.0) + monkeypatch.setenv(RANK_CRASH_KILL_GRACE_ENV, "30") + + # 29.5s of the 30s grace has already been burned by the watchdog. + assert hard_kill_on_rank_crash(world_size=2, deadline=1000.5) is True + # Exactly one 0.5s sleep, then exactly one kill -- never a second grace. + assert order.count(("sleep", 0.5)) == 1, order + assert not any(s == ("sleep", 30.0) for s in order), order + _assert_slept_then_killed(order, 0.5) + + +def test_kill_fires_immediately_when_deadline_already_passed(monkeypatch): + """A deadline in the past must fire the kill now, and never sleep negative. + + time.sleep of a negative duration raises into hard_kill_on_rank_crash's + blanket except, which would return False and silently skip the kill -- + exactly in the case the watchdog exists for (cleanup outlasted the grace). + """ + order = [] + monkeypatch.setattr(hang_detector_module, "propagate_hard_kill", lambda: order.append("kill")) + monkeypatch.setattr(hang_detector_module.time, "sleep", lambda s: order.append(("sleep", s))) + monkeypatch.setenv(RANK_CRASH_KILL_GRACE_ENV, "30") + + t0 = time.monotonic() + assert hard_kill_on_rank_crash(world_size=2, deadline=t0 - 100.0) is True + assert order == ["kill"], order + assert not [s for s in order if isinstance(s, tuple) and s[1] < 0], order + + +def test_wait_out_kill_grace_never_sleeps_negative(monkeypatch): + """The `remaining > 0` guard, not the deadline clamp, is what protects here.""" + slept = [] + monkeypatch.setattr(hang_detector_module.time, "sleep", lambda s: slept.append(s)) + assert hang_detector_module._wait_out_kill_grace(-100.0, None) is True + assert slept == [] + # The cancellable path must also return promptly, not wait forever. + assert hang_detector_module._wait_out_kill_grace(-100.0, threading.Event()) is True + + +# -------------------------------------------------------------------------- +# 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, is_shutdown=False): + # 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 + # is_shutdown must NOT influence the crash decision -- _handle_errors sets + # it rank-locally on a fatal error while peers are told nothing. + ex.is_shutdown = is_shutdown + ex._event_loop_completed = False + # The real __init__ creates this; __new__ does not run it. The wrapper + # reads it on every crash path, so a bare executor without it turns a + # wiring test into an AttributeError. + ex._event_loop_error_delivered = threading.Event() + return ex + + +class _FakeWatchdog: + """Stand-in for RankCrashKillWatchdog that records cancellation. + + ``cancelled`` is a read-only property, matching the real class: a wiring + change that assigned to it would pass against a plain attribute here and + raise AttributeError in production. + """ + + def __init__(self, events, world_size): + self._events = events + self._cancelled = False + self.deadline = 1234.5 + events.append(("watchdog", world_size)) + + @property + def cancelled(self): + return self._cancelled + + def cancel(self): + self._cancelled = True + self._events.append("cancel") + + +def _stub_kill_paths(pe, monkeypatch, events, arm_watchdog=True, seen=None): + # ``error_delivered`` is keyword-only with NO default on both stubs: if the + # wiring that threads the delivery gate through is ever dropped, these + # raise TypeError instead of silently accepting the pre-gate signature. + # Pass ``seen`` to capture the objects actually handed over. + def _kill(world_size, deadline=None, *, error_delivered): + if seen is not None: + seen.append(("kill", error_delivered)) + events.append(("kill", world_size, deadline)) + + monkeypatch.setattr(pe, "hard_kill_on_rank_crash", _kill) + watchdogs = [] + + def _start(world_size, *, error_delivered): + if seen is not None: + seen.append(("watchdog", error_delivered)) + if not arm_watchdog: + events.append(("watchdog", world_size)) + return None + wd = _FakeWatchdog(events, world_size) + watchdogs.append(wd) + return wd + + monkeypatch.setattr(pe, "start_rank_crash_kill_watchdog", _start) + return watchdogs + + +def test_event_loop_wrapper_kills_world_on_crash(monkeypatch): + """A genuine mid-loop crash (is_shutdown still False) must kill the world.""" + from tensorrt_llm._torch.pyexecutor import py_executor as pe + + events = [] + watchdogs = _stub_kill_paths(pe, monkeypatch, events) + ex = _bare_executor(pe, monkeypatch, world_size=4, is_shutdown=False) + 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. Once cleanup returns, the + # watchdog is disarmed and the kill is carried inline on the watchdog's + # ORIGINAL deadline, so only one timer is ever live. + assert events == [("watchdog", 4), "cleanup", "cancel", ("kill", 4, 1234.5)] + assert watchdogs[0].cancelled is True + assert isinstance(ex._event_loop_error, ValueError) + + +def test_event_loop_wrapper_hands_both_kill_paths_this_executors_gate(monkeypatch): + """Both kill paths must receive THIS executor's delivery gate. + + Handing over a fresh Event, or a different executor's, would read as + "the error never reached the client" and kill the world even on the + path the grace exists to protect. + """ + from tensorrt_llm._torch.pyexecutor import py_executor as pe + + events, seen = [], [] + _stub_kill_paths(pe, monkeypatch, events, seen=seen) + ex = _bare_executor(pe, monkeypatch, world_size=4, is_shutdown=False) + 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() + + assert [kind for kind, _ in seen] == ["watchdog", "kill"] + assert all(gate is ex._event_loop_error_delivered for _, gate in seen) + + +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, is_shutdown=False) + + 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", "cancel", ("kill", 4, 1234.5)] + # The original loop error stays reachable for rank-local consumers. + assert isinstance(ex._event_loop_error, ValueError) + + +def test_event_loop_wrapper_kills_world_when_watchdog_cannot_arm(monkeypatch): + """A watchdog that fails to start must not silently drop the escalation.""" + from tensorrt_llm._torch.pyexecutor import py_executor as pe + + events = [] + _stub_kill_paths(pe, monkeypatch, events, arm_watchdog=False) + ex = _bare_executor(pe, monkeypatch, world_size=4, is_shutdown=False) + 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() + + assert events == [("watchdog", 4), "cleanup", ("kill", 4, None)] + + +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"] + + +# -------------------------------------------------------------------------- +# The kill must stay scoped to crashes that actually strand peers, and the +# only signal that says so is _event_loop_completed -- set at the loops' +# normal-exit `break` sites and nowhere else. is_shutdown does NOT mean +# "peers were told": _handle_errors flips it rank-locally on a fatal error. +# -------------------------------------------------------------------------- + + +def test_event_loop_wrapper_no_kill_when_loop_raises_after_completing(monkeypatch): + """A raise after the loop's normal-exit break is a teardown error, not a crash.""" + 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 late_raise(): + # The loop hit its normal-exit `break` and drained all work, then + # something raised on the way out (e.g. a context manager's __exit__). + ex._event_loop_completed = True + raise RuntimeError("teardown hiccup") + + ex.event_loop = late_raise + + with pytest.raises(RuntimeError, match="teardown hiccup"): + ex._event_loop_wrapper() + + # Logged and re-raised, but no watchdog and no kill: peers are not stranded. + assert events == ["cleanup"] + assert isinstance(ex._event_loop_error, RuntimeError) + + +def test_event_loop_wrapper_kills_world_on_rank_local_fatal(monkeypatch): + """REGRESSION: a rank-local CUDA fatal sets is_shutdown but strands peers. + + _handle_errors classifies a device-side fault as immediate_fatal, sets + is_shutdown=True on THIS rank and enqueues a shutdown into THIS process's + own queue -- peers are told nothing and keep waiting in their collective. + An exception raised after that point (e.g. the unguarded + guided_decoder.execute(batch_outputs['logits']) on a None batch_outputs) + must still hard-kill the world. Keying the decision off is_shutdown + silently disabled the kill for this, the feature's most common trigger. + """ + 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 cuda_fatal_then_crash(): + ex.is_shutdown = True # what _handle_errors does, rank-locally + assert ex._event_loop_completed is False # the loop never terminated + raise TypeError("'NoneType' object is not subscriptable") + + ex.event_loop = cuda_fatal_then_crash + + with pytest.raises(TypeError): + ex._event_loop_wrapper() + + assert events == [("watchdog", 4), "cleanup", "cancel", ("kill", 4, 1234.5)] + + +def test_event_loop_wrapper_kills_world_when_loop_never_started(monkeypatch): + """A failure before the loop runs strands peers just as surely as one inside it.""" + from tensorrt_llm._torch.pyexecutor import py_executor as pe + + events = [] + _stub_kill_paths(pe, monkeypatch, events) + + @contextlib.contextmanager + def failing_enter(**_kwargs): + raise RuntimeError("profiler setup failed") + yield # pragma: no cover + + ex = _bare_executor(pe, monkeypatch, world_size=4) + monkeypatch.setattr(pe, "host_profiler_context", lambda enable: failing_enter()) + ex._executor_loop_cleanup = lambda: events.append("cleanup") + ex.event_loop = lambda: pytest.fail("event_loop must not be reached") + + with pytest.raises(RuntimeError, match="profiler setup failed"): + ex._event_loop_wrapper() + + assert events == [("watchdog", 4), "cleanup", "cancel", ("kill", 4, 1234.5)] + + +def test_event_loop_wrapper_no_kill_when_enclosing_context_manager_raises(monkeypatch): + """Teardown of the host-profiler / GC context managers after a completed loop. + + They wrap event_loop() but are not part of it; a failure while unwinding + them once the loop has completed leaves no peer waiting on this rank. + """ + from tensorrt_llm._torch.pyexecutor import py_executor as pe + + events = [] + _stub_kill_paths(pe, monkeypatch, events) + + @contextlib.contextmanager + def exploding_ctx(**_kwargs): + yield + raise RuntimeError("profiler teardown failed") + + ex = _bare_executor(pe, monkeypatch, world_size=4) + monkeypatch.setattr(pe, "host_profiler_context", lambda enable: exploding_ctx()) + ex._executor_loop_cleanup = lambda: events.append("cleanup") + + def completed_loop(): + ex._event_loop_completed = True + + ex.event_loop = completed_loop + + with pytest.raises(RuntimeError, match="profiler teardown failed"): + ex._event_loop_wrapper() + + assert events == ["cleanup"] + + +# -------------------------------------------------------------------------- +# The sentinel is only trustworthy if the real loops actually set it. Assert +# against the shipped source so a new normal-exit path (or a moved break) +# cannot silently make every clean shutdown look like a peer-stranding crash. +# -------------------------------------------------------------------------- + + +def _executor_loop_ast_nodes(): + import ast + import inspect + + from tensorrt_llm._torch.pyexecutor import py_executor as pe + + tree = ast.parse(inspect.getsource(pe)) + wanted = {"_executor_loop", "_executor_loop_pp", "_executor_loop_overlap"} + return { + node.name: node + for node in ast.walk(tree) + if isinstance(node, ast.FunctionDef) and node.name in wanted + } + + +def _outer_while(fn): + """The `while True:` that IS the event loop. + + Selected by its `True` test, not by walk order: _executor_loop_pp contains + three other `while`s, one of them (the Stage-5 drain) a SIBLING of the + event loop, so relying on ast.walk ordering would silently point the guard + at the wrong loop if the body were ever reordered. + """ + import ast + + candidates = [ + node + for node in ast.walk(fn) + if isinstance(node, ast.While) + and isinstance(node.test, ast.Constant) + and node.test.value is True + ] + assert len(candidates) == 1, ( + f"{fn.name}: expected exactly one `while True:` (the event loop), " + f"found {len(candidates)} at lines {[c.lineno for c in candidates]}" + ) + return candidates[0] + + +def _loop_terminating_breaks(loop): + """(block, index, break_node) for every break that exits ``loop`` itself. + + Recurses through if/try/with, but stops at nested for/while/def: a break + inside those binds to the inner construct, not to the event loop. + """ + import ast + + found = [] + + def visit(block): + for i, stmt in enumerate(block): + if isinstance(stmt, ast.Break): + found.append((block, i, stmt)) + elif isinstance( + stmt, (ast.For, ast.AsyncFor, ast.While, ast.FunctionDef, ast.AsyncFunctionDef) + ): + continue # binds to the inner construct + else: + for field in ("body", "orelse", "finalbody", "handlers"): + inner = getattr(stmt, field, None) + if isinstance(inner, list): + if field == "handlers": + for h in inner: + visit(h.body) + else: + visit(inner) + + visit(loop.body) + return found + + +def _sets_sentinel_true(stmt): + """Exactly `self._event_loop_completed = True` -- object and value both checked.""" + import ast + + if not isinstance(stmt, ast.Assign): + return False + if not (isinstance(stmt.value, ast.Constant) and stmt.value.value is True): + return False + return any( + isinstance(t, ast.Attribute) + and t.attr == "_event_loop_completed" + and isinstance(t.value, ast.Name) + and t.value.id == "self" + for t in stmt.targets + ) + + +def test_loop_terminating_break_sets_the_completion_sentinel(): + """Only the break that exits the OUTER `while True` terminates the event loop. + + Deliberately not "every break": these loops contain inner `for` loops, and + an inner break does not end the event loop. Demanding the sentinel there + would instruct a contributor to set it while the loop is still running, + which makes every later rank-local crash look like a clean shutdown and + silently disables the kill -- the same class of bug this predicate has + already regressed into twice. + """ + + loops = _executor_loop_ast_nodes() + assert set(loops) == {"_executor_loop", "_executor_loop_pp", "_executor_loop_overlap"}, ( + f"executor loops renamed or removed: {sorted(loops)}" + ) + + for name, fn in loops.items(): + outer = _outer_while(fn) + assert outer is not None, f"{name}: no `while` loop found -- did the loop shape change?" + + terminating = _loop_terminating_breaks(outer) + assert len(terminating) == 1, ( + f"{name}: expected exactly 1 loop-terminating break, found " + f"{len(terminating)} at lines {[b.lineno for _, _, b in terminating]}. " + "A new normal-exit path must also set self._event_loop_completed = True." + ) + + block, idx, brk = terminating[0] + prev = block[idx - 1] if idx else None + assert _sets_sentinel_true(prev), ( + f"{name}: the loop-terminating `break` at line {brk.lineno} is not " + "preceded by `self._event_loop_completed = True`. Without it " + "_event_loop_wrapper treats a clean shutdown as a peer-stranding " + "crash and SIGKILLs the job. (Inner-loop breaks must NOT set it.)" + ) + + +def test_completion_sentinel_is_reset_per_event_loop_run(): + """The reset in _event_loop_wrapper is load-bearing, not redundant with __init__. + + PyExecutor outlives a single loop run; without the reset a second run + starts with the sentinel left True by the first, so a genuine crash in it + is misread as a clean shutdown and no kill is armed. + """ + import inspect + + from tensorrt_llm._torch.pyexecutor import py_executor as pe + + assert "self._event_loop_completed = False" in inspect.getsource( + pe.PyExecutor._event_loop_wrapper + ) + + +def test_second_loop_run_still_kills_after_a_clean_first_run(monkeypatch): + """Behavioral guard for the reset above: run clean, then crash, on ONE executor.""" + 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") + + # Run 1: reaches the normal-exit break, leaving the sentinel True. + def clean(): + ex._event_loop_completed = True + + ex.event_loop = clean + ex._event_loop_wrapper() + assert events == ["cleanup"] + assert ex._event_loop_completed is True + + # Run 2 on the SAME executor: a genuine crash must still arm the kill. + events.clear() + ex.event_loop = lambda: (_ for _ in ()).throw(ValueError("boom")) + with pytest.raises(ValueError, match="boom"): + ex._event_loop_wrapper() + + assert events == [("watchdog", 4), "cleanup", "cancel", ("kill", 4, 1234.5)] + + +# --------------------------------------------------------------------------- +# The delivery gate (review: symmetric crashes must not become exit 137). +# --------------------------------------------------------------------------- + + +def test_kill_is_skipped_once_the_error_reached_the_client(monkeypatch): + """A reportable crash must not be converted into a bare exit 137. + + `crashed` means "the loop raised before its break", which is broader than + "peers are stranded". In a symmetric crash every rank raises, nobody is + stranded, and every rank arms this kill. If the stashed error already + surfaced to the client the failure is diagnosable, so killing the world + only destroys N tracebacks. + """ + calls = [] + monkeypatch.setattr( + hang_detector_module, "propagate_hard_kill", lambda *a, **k: calls.append(1) + ) + monkeypatch.setenv(RANK_CRASH_KILL_GRACE_ENV, "0") + + delivered = threading.Event() + delivered.set() + fired = hard_kill_on_rank_crash(4, error_delivered=delivered) + + assert fired is False, "kill fired despite the error having been delivered" + assert calls == [], "propagate_hard_kill must not run once the error is reportable" + + +def test_kill_still_fires_when_nothing_consumed_the_error(monkeypatch): + """The stranded-peer case is unchanged: nothing consumes it, so kill.""" + calls = [] + monkeypatch.setattr( + hang_detector_module, "propagate_hard_kill", lambda *a, **k: calls.append(1) + ) + monkeypatch.setenv(RANK_CRASH_KILL_GRACE_ENV, "0") + + fired = hard_kill_on_rank_crash(4, error_delivered=threading.Event()) + + assert fired is True + assert calls == [1] + + +def test_kill_fires_when_no_delivery_event_is_supplied(monkeypatch): + """Back-compat: callers that pass nothing get the old behaviour.""" + calls = [] + monkeypatch.setattr( + hang_detector_module, "propagate_hard_kill", lambda *a, **k: calls.append(1) + ) + monkeypatch.setenv(RANK_CRASH_KILL_GRACE_ENV, "0") + + assert hard_kill_on_rank_crash(4) is True + assert calls == [1] + + +def test_delivery_is_checked_after_the_grace_not_before(monkeypatch): + """The check must come after the wait, else it defeats its own purpose. + + The grace exists so the error can reach the client. Sampling the flag + before waiting would read it while it is still False and kill anyway. + """ + calls = [] + monkeypatch.setattr( + hang_detector_module, "propagate_hard_kill", lambda *a, **k: calls.append(1) + ) + monkeypatch.setenv(RANK_CRASH_KILL_GRACE_ENV, "0.5") + + delivered = threading.Event() + + def deliver_during_grace(): + time.sleep(0.15) + delivered.set() + + t = threading.Thread(target=deliver_during_grace, daemon=True) + t.start() + fired = hard_kill_on_rank_crash(4, error_delivered=delivered) + t.join(timeout=5) + + assert fired is False, "the flag was set during the grace window; the kill must observe it" + assert calls == [] + + +def test_watchdog_threads_the_delivery_event_through(monkeypatch): + calls = [] + monkeypatch.setattr( + hang_detector_module, "propagate_hard_kill", lambda *a, **k: calls.append(1) + ) + monkeypatch.setenv(RANK_CRASH_KILL_GRACE_ENV, "0") + + delivered = threading.Event() + delivered.set() + wd = start_rank_crash_kill_watchdog(4, error_delivered=delivered) + assert wd is not None + wd.join(timeout=5) + + assert calls == [], "watchdog killed despite a delivered error" + + +# --------------------------------------------------------------------------- +# Real 2-rank MPI: the kill and the gate, with propagate_hard_kill NOT mocked. +# +# Every other kill-path test in this file monkeypatches propagate_hard_kill, +# so none of them exercises a real MPI_Abort (raised in review by @BowenFu). +# These two do: they launch a real 2-rank MPI job and assert on the exit +# status of the whole job. +# +# Scope, stated honestly: this proves the KILL MECHANISM and the delivery +# gate over a real communicator. It is not a full 2-rank LLM crash -- there +# is no engine here -- so it does not by itself prove the end-to-end claim +# that a client sees the original exception. It does close the "nothing +# exercises a real MPI_Abort" gap. +# --------------------------------------------------------------------------- + +_MPI_2RANK_SCRIPT = """ +import os, sys, time +from mpi4py import MPI +from tensorrt_llm._torch.pyexecutor.hang_detector import hard_kill_on_rank_crash + +comm = MPI.COMM_WORLD +comm.Barrier() # both ranks up, imports done +# Printed only once imports and MPI init have succeeded. The assertions +# require it, so a setup failure (bad import, no MPI) cannot masquerade as +# a successful abort just by exiting non-zero. +if comm.Get_rank() == 0: + print("RANK0_READY", flush=True) + +if comm.Get_rank() == 0: + import threading + delivered = threading.Event() + if os.environ["DELIVERED"] == "1": + delivered.set() + hard_kill_on_rank_crash(comm.Get_size(), error_delivered=delivered) + # Only reached when the kill is skipped. + print("RANK0_SURVIVED", flush=True) +else: + # A peer that would otherwise sit in a collective forever. + time.sleep(20) + print("RANK1_SURVIVED", flush=True) + +comm.Barrier() +sys.exit(0) +""" + + +def _run_two_rank(delivered: str): + env = { + **os.environ, + "DELIVERED": delivered, + hang_detector_module.RANK_CRASH_KILL_GRACE_ENV: "0", + } + return subprocess.run( + ["mpirun", "--allow-run-as-root", "-n", "2", sys.executable, "-c", _MPI_2RANK_SCRIPT], + env=env, + timeout=600, + capture_output=True, + ) + + +@pytest.mark.skipif(shutil.which("mpirun") is None, reason="mpirun not available") +def test_real_mpi_abort_takes_down_both_ranks(): + """Undelivered crash: the abort must reach the peer, not just rank 0. + + Cross-rank propagation is the load-bearing part of the whole feature. If + MPI_Abort only killed rank 0, the peer would still burn to its own + HangDetector -- exactly the failure this exists to prevent. + """ + proc = _run_two_rank(delivered="0") + out = (proc.stdout + proc.stderr).decode(errors="replace") + + assert "RANK0_READY" in out, ( + f"the job never reached the kill call -- this is a setup failure, not " + f"an abort, and must not be read as a pass; out={out[-1500:]}" + ) + assert proc.returncode != 0, f"job survived an undelivered crash kill; out={out[-800:]}" + assert "RANK1_SURVIVED" not in out, ( + f"peer rank outlived the abort -- propagation failed; out={out[-800:]}" + ) + + +@pytest.mark.skipif(shutil.which("mpirun") is None, reason="mpirun not available") +def test_real_mpi_job_survives_when_the_error_was_delivered(): + """Delivered crash: no abort, so both ranks run to completion. + + This is the review point -- a symmetric crash whose error already reached + the client must not have its tracebacks replaced by exit 137. + """ + proc = _run_two_rank(delivered="1") + out = (proc.stdout + proc.stderr).decode(errors="replace") + + assert "RANK0_READY" in out, ( + f"the job never reached the kill call -- setup failure; out={out[-1500:]}" + ) + assert proc.returncode == 0, f"job died despite a delivered error; out={out[-800:]}" + assert "RANK0_SURVIVED" in out, f"rank 0 was killed anyway; out={out[-800:]}" + assert "RANK1_SURVIVED" in out, f"peer was killed anyway; out={out[-800:]}" diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index c9c05b39137f..820b9702e8e9 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -1111,6 +1111,10 @@ def __init__(self): self.responses = {} self.is_shutdown = False self._event_loop_error = None + # Set when the stashed error is handed to a caller: that is what tells + # the rank-crash kill the crash was already reported and the world does + # not need tearing down. + self._event_loop_error_delivered = threading.Event() # Bind the real production method so the test exercises real code. _await_single_response = PyExecutor._await_single_response @@ -1149,6 +1153,10 @@ def test_raises_on_shutdown_with_event_loop_error(self): with pytest.raises(RuntimeError, match="Event loop terminated"): stub._await_single_response(id=42, timeout=1.0) + # The caller now holds the original error, so the rank-crash kill must + # stand down: this is the signal it waits out its grace for. + assert stub._event_loop_error_delivered.is_set() + def test_raises_on_shutdown_without_event_loop_error(self): """Shutdown without a stored error still raises rather than blocking — distinguishes "shutdown" from "timed out without shutdown".""" @@ -1158,6 +1166,10 @@ def test_raises_on_shutdown_without_event_loop_error(self): with pytest.raises(RuntimeError, match="Event loop shut down"): stub._await_single_response(id=42, timeout=1.0) + # Nothing was delivered -- there was no error to deliver. Leaving the + # gate clear keeps the kill armed, which is correct here. + assert not stub._event_loop_error_delivered.is_set() + def test_returns_empty_on_timeout(self): """Pre-fix behaviour: a bare timeout (no shutdown, no response) used to KeyError. The fix returns an empty list to match the documented 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