From 9c7cd069c386be5f5367e6e5c5c0d59bb0a6b373 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Mon, 10 Aug 2026 02:49:51 +0000 Subject: [PATCH] [TRTLLM-13409][feat] hard-kill all ranks when one rank's executor loop crashes When a rank's executor loop dies on an exception, that rank stops participating in collectives but nothing tells its peers. Every peer blocks in its next collective until its own HangDetector fires 300s later, so the whole multi-GPU session burns that long for an error that was already known. The A4 AutoDeploy catches are this signature: peers crash, the survivor wedges in ADP until the 300s backstop. A rank whose executor loop crashes now tears the job down (MPI_Abort, falling back to self-SIGKILL) after a grace period. Single-rank worlds are exempt -- there are no peers to strand -- and the kill helper never raises, since it runs in a `finally` where an exception would mask the loop's original error. A cancellable RankCrashKillWatchdog is armed before cleanup so the kill stays reachable when cleanup itself blocks or raises, and the original fire deadline survives handover between waiters. The grace exists so cleaner paths can win the race: if the crash is surfaced to a client first, the caller gets the real traceback and the kill stands down rather than replacing it with a bare exit 137. That gate is set only on verified delivery -- the helper must be on the single-process path AND a readable queue must have taken the error -- because on the proxy/IPC path the broadcast cannot reach the client at all. Setting it there would disarm the kill while the peers were still stranded, which is the case this change exists for, in the default spawned-worker deployment. The stand-down is not guaranteed for symmetric crashes: only rank 0 can set the gate, so subordinate kills still fire at crash+grace and the "N tracebacks" outcome holds only if each subordinate exits within the grace. destroy_process_group() on a wedged NCCL communicator can exceed it. That is deliberate -- the timer doubles as protection against wedged teardown. This enables world-kill-on-crash by default for every multi-rank run. TLLM_RANK_CRASH_HARD_KILL_GRACE tunes the grace; any negative value (e.g. -1) disables the kill entirely and restores the previous hang-detector behavior. Unset, unparsable and non-finite values use the 10s default -- nan in particular would otherwise slip past the `grace < 0` check and collapse to a zero grace, and inf would produce a watchdog that never fires. Both the default change and the escape hatch are documented in docs/source/developer-guide/overview.md. Tests: test_hang_detector_kill.py (l0_sanity_check) covers the exemption, grace timing and ordering, disable/invalid/non-finite env handling, non-raising behavior, watchdog arm/cancel/deadline-handover and the crash-vs-clean-exit wiring, including two real 2-rank MPI_Abort tests; test_event_loop_error_broadcast.py (l0_cpu) pins the delivery gate, including that it stays clear on ipc_batched. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- docs/source/developer-guide/overview.md | 29 + .../_torch/pyexecutor/hang_detector.py | 245 +++++ tensorrt_llm/_torch/pyexecutor/py_executor.py | 100 +- tensorrt_llm/executor/base_worker.py | 51 +- .../executor/test_hang_detector_kill.py | 916 +++++++++++++++++- .../_torch/executor/test_py_executor.py | 12 + .../test_event_loop_error_broadcast.py | 75 +- 7 files changed, 1420 insertions(+), 8 deletions(-) diff --git a/docs/source/developer-guide/overview.md b/docs/source/developer-guide/overview.md index d8fe31631612..104a6957e330 100644 --- a/docs/source/developer-guide/overview.md +++ b/docs/source/developer-guide/overview.md @@ -122,3 +122,32 @@ export TLLM_LOG_LEVEL_BY_MODULE="debug:_torch,runtime;info:serve" ``` This example sets the global level to `warning` but enables `debug` output for `_torch` and `runtime` modules, and `info` for `serve`. Valid levels: `trace`, `debug`, `verbose`, `info`, `warning`, `error`, `internal_error`. + +## Multi-Rank Crash Handling + +When one rank's executor loop dies in a multi-rank job, the peer ranks have +nothing to fail on: they sit in a collective waiting for a rank that will never +arrive. Left alone they hold every GPU in the job until an external wall-clock +kill. + +To avoid that, a rank whose executor loop crashes tears the whole job down +(`MPI_Abort`, falling back to self-`SIGKILL`) after a short grace period. The +grace exists so cleaner paths can win the race: if the crash is reported to a +client first — the caller gets the real traceback — the kill stands down, so a +symmetric crash still ends in N tracebacks rather than a bare exit 137. + +**This is on by default for multi-rank runs.** Single-rank runs are unaffected; +there are no peers to strand. + +| Variable | Default | Meaning | +|---|---|---| +| `TLLM_RANK_CRASH_HARD_KILL_GRACE` | `10` | Seconds to wait before killing the job. `0` kills immediately. Any negative value (e.g. `-1`) **disables** the kill entirely — peers then fall back to the hang detector. Unparsable and non-finite values (`nan`, `inf`) fall back to the default. | + +```bash +# Opt out: let the hang detector handle it instead (slower, but no world kill). +export TLLM_RANK_CRASH_HARD_KILL_GRACE=-1 +``` + +Use the escape hatch when you would rather have peer ranks reach their own +timeouts — for example when attaching a debugger to a surviving rank, or when a +harness collects per-rank state that a job-wide abort would destroy. diff --git a/tensorrt_llm/_torch/pyexecutor/hang_detector.py b/tensorrt_llm/_torch/pyexecutor/hang_detector.py index f6dde7c58959..9773ff324849 100644 --- a/tensorrt_llm/_torch/pyexecutor/hang_detector.py +++ b/tensorrt_llm/_torch/pyexecutor/hang_detector.py @@ -13,10 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. import asyncio +import math import os import signal import sys import threading +import time from contextlib import contextmanager from typing import Callable, Optional @@ -26,6 +28,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 +51,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 +95,236 @@ 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 + if not math.isfinite(grace): + # nan slips past `grace < 0` (every nan comparison is False) and then + # collapses to a ZERO grace downstream: max(0.0, nan) returns 0.0, and + # the direct path's `remaining > 0` guard is False too. That kills + # instantly, destroying the very window the grace exists to provide. + # inf is the mirror case: a watchdog thread that never fires, silently + # equivalent to -1 but costing a live thread. Reject both. + _best_effort_log_error( + f"Non-finite {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 cb8bff1a5bea..2490dea8f3ea 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -67,7 +67,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, @@ -652,6 +653,21 @@ 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. + # + # Only rank 0 can ever set this: both delivery sites need a client + # consumer (_await_single_response, and AwaitResponseHelper on the + # single-process path), and subordinate ranks block in wait_shutdown() + # with no response thread. So in a SYMMETRIC crash the subordinates' + # kills still fire at crash+grace, and the "N clean tracebacks instead + # of exit 137" outcome only holds if every subordinate finishes cleanup + # and exits within the grace -- destroy_process_group() on a wedged + # NCCL communicator can exceed it. That is the intended tradeoff: the + # timer doubles as protection against wedged teardown. + self._event_loop_error_delivered = threading.Event() # kv cache events self.kv_cache_manager = self.resource_manager.resource_managers.get( @@ -860,6 +876,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() @@ -1217,6 +1242,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 @@ -1226,6 +1253,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 @@ -1237,8 +1278,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: @@ -2583,6 +2667,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() @@ -2884,6 +2969,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() @@ -4066,6 +4157,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( @@ -4540,6 +4632,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( @@ -7429,6 +7522,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..51845c9425e3 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -1087,7 +1087,14 @@ def __init__(self, worker: "BaseWorker"): # The error responses when submit request failed will be put here self.temp_error_responses = Queue() - def responses_handler(self, responses: List[tllm.Response]): + def _resolve_handler_kind(self) -> "AwaitResponseHelper.HandlerKind": + """Determine (and memoise) which side of the IPC boundary we are on. + + Split out of ``responses_handler`` so the error path can ask the same + question without having handled a response batch first — a crash + during the very first ``await_responses`` leaves ``handler_kind`` + ``unknown`` otherwise. + """ HandlerKind = AwaitResponseHelper.HandlerKind if self.handler_kind is HandlerKind.unknown: @@ -1105,6 +1112,12 @@ def responses_handler(self, responses: List[tllm.Response]): logger_debug(f"creating await_response helper for IPC\n", color="yellow") self.handler_kind = HandlerKind.ipc_batched + return self.handler_kind + + def responses_handler(self, responses: List[tllm.Response]): + HandlerKind = AwaitResponseHelper.HandlerKind + + self._resolve_handler_kind() match self.handler_kind: case HandlerKind.single_process_worker: @@ -1152,6 +1165,8 @@ 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: + # _broadcast_event_loop_error owns the delivery gate: it is the + # only place that knows whether a client was actually woken. return self._broadcast_event_loop_error(error) return True @@ -1172,8 +1187,21 @@ def _broadcast_event_loop_error(self, error: BaseException) -> bool: results on a different side of the boundary and would need a separate poison-pill on ``self.worker.result_queue``; that is left as a follow-up consistent with the PyExecutor-side fix. + + Because of that scope, this method also owns the rank-crash kill's + delivery gate. The gate may only be set when a client verifiably + woke: on ``ipc_batched`` the queues written below have no reader + (responses travel via ``handle_for_ipc_batched``), so setting it + there would stand the kill down while the peer ranks are still + stranded — the case the kill exists for, in the default spawned- + worker deployment. When delivery cannot be proven the gate stays + clear and the kill fires, which is the safe direction: a spurious + world-kill costs a traceback, a missed one costs the job. """ error_msg = f"Event loop terminated with error: {error}" + can_reach_client = ( + self._resolve_handler_kind() + is AwaitResponseHelper.HandlerKind.single_process_worker) pending_client_ids = list(self.worker._results.keys()) if not pending_client_ids: logger.error( @@ -1186,6 +1214,9 @@ def _broadcast_event_loop_error(self, error: BaseException) -> bool: event_loop = None async_queues: List[_SyncQueue] = [] + # Counts queues a caller can actually read from. A _SyncQueue is only + # readable once notify_many() has run, so those are counted there. + woken = 0 for client_id in pending_client_ids: try: queue = self.worker.return_queue(client_id) @@ -1203,6 +1234,7 @@ def _broadcast_event_loop_error(self, error: BaseException) -> bool: event_loop = event_loop or queue.loop else: queue.put(err_resp) + woken += 1 except Exception as put_error: logger.error(f"Failed to push ErrorResponse for client_id=" f"{client_id}: {put_error}") @@ -1212,11 +1244,28 @@ def _broadcast_event_loop_error(self, error: BaseException) -> bool: if async_queues: try: _SyncQueue.notify_many(event_loop, async_queues) + woken += len(async_queues) except Exception as notify_error: logger.error( f"Failed to notify async queues on event-loop error: " f"{notify_error}") + if woken and can_reach_client: + # A client is now holding the real error, so the crash is + # reportable without killing the world: 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() + elif not can_reach_client: + logger.error( + "Event-loop error broadcast cannot reach the client on the " + "IPC/proxy path; leaving the rank-crash hard kill armed so " + "peer ranks are not stranded.") + return False def handle_for_worker(self, responses: List[tllm.Response]) -> None: diff --git a/tests/unittest/_torch/executor/test_hang_detector_kill.py b/tests/unittest/_torch/executor/test_hang_detector_kill.py index 0962df441cc6..210c3bbbda82 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,906 @@ 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) + + +@pytest.mark.parametrize("raw", ["nan", "NaN", "inf", "-inf", "Infinity"]) +def test_rank_crash_kill_non_finite_grace_uses_default(monkeypatch, raw): + """float() accepts nan/inf, but neither survives the arithmetic downstream. + + ``nan`` slips past the ``grace < 0`` check (all nan comparisons are False) + and then collapses to a ZERO grace -- ``max(0.0, nan)`` is ``0.0`` and + ``remaining > 0`` is False -- so the kill fires instantly and destroys the + window the grace exists to give the traceback. ``inf`` is the mirror case: + a watchdog that never fires, silently equivalent to ``-1`` but costing a + live thread. Both must fall back to the documented 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, raw) + 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 fa8441426dbf..4f84a7f6908e 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -1387,6 +1387,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 @@ -1425,6 +1429,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".""" @@ -1434,6 +1442,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_event_loop_error_broadcast.py b/tests/unittest/executor/test_event_loop_error_broadcast.py index f5b72d22e3f6..8e061b11930d 100644 --- a/tests/unittest/executor/test_event_loop_error_broadcast.py +++ b/tests/unittest/executor/test_event_loop_error_broadcast.py @@ -12,6 +12,7 @@ import datetime import queue as _stdlib_queue +import threading import pytest @@ -36,6 +37,9 @@ def __init__( self._event_loop_error = event_loop_error self.is_shutdown = is_shutdown self.calls = 0 + # The rank-crash kill's delivery gate. Real PyExecutor creates this in + # __init__; the helper only ever sets it, never reads it. + self._event_loop_error_delivered = threading.Event() def await_responses(self, timeout: datetime.timedelta): self.calls += 1 @@ -54,11 +58,13 @@ def __init__(self): class _WorkerStub: """Stub for BaseWorker exposing only the attributes the helper touches.""" - def __init__(self, engine, num_pending: int = 1): + def __init__(self, engine, num_pending: int = 1, ipc: bool = False): self.engine = engine self._results = {cid: _ResultStub() for cid in range(1, num_pending + 1)} self.popped = [] - self.result_queue = None + # A non-None result_queue is what makes the helper resolve to + # ipc_batched -- i.e. the proxy/spawned-worker deployment. + self.result_queue = object() if ipc else None self.postproc_queues = None # responses_handler() reads this unguarded (base_worker.py); BaseWorker # sets it in __init__, which this stub bypasses, so seed it to None. @@ -78,9 +84,9 @@ def _pop_result(self, client_id: int): self._results.pop(client_id, None) -def _make_helper(engine, num_pending: int = 1): +def _make_helper(engine, num_pending: int = 1, ipc: bool = False): helper = AwaitResponseHelper.__new__(AwaitResponseHelper) - helper.worker = _WorkerStub(engine, num_pending=num_pending) + helper.worker = _WorkerStub(engine, num_pending=num_pending, ipc=ipc) helper.handler_kind = AwaitResponseHelper.HandlerKind.unknown helper.enable_postprocprocess_parallel = False helper.temp_error_responses = _stdlib_queue.Queue() @@ -162,3 +168,64 @@ def test_broadcast_helper_idempotent_via_pop(self): assert sorted(helper.worker.popped) == [1, 2] # second time around: nothing left to wake. assert helper._broadcast_event_loop_error(original) is False + + +class TestEventLoopErrorDeliveryGate: + """The gate that stands the rank-crash hard kill down. + + Setting it means "a client is holding the real error, so killing the + world would only replace a traceback with exit 137". It may therefore + only be set when a client verifiably woke. Setting it optimistically on + the proxy/IPC path -- the default when the LLM spawns MPI workers -- + disarms the 10s kill while the peer ranks are still stranded, regressing + exactly the case the kill exists for back to the 300s HangDetector. + """ + + def test_gate_set_when_client_woken_in_single_process_mode(self): + original = RuntimeError("KV cache OOM") + engine = _EngineStub(event_loop_error=original, is_shutdown=True) + helper = _make_helper(engine, num_pending=2, ipc=False) + + assert helper(timeout=0.01) is False + assert helper.handler_kind is AwaitResponseHelper.HandlerKind.single_process_worker + assert engine._event_loop_error_delivered.is_set() + + def test_gate_stays_clear_on_ipc_batched_path(self): + """The regression this class exists for. + + On ``ipc_batched`` the worker-side ``_results`` queues written by the + broadcast have no reader -- responses travel via + ``handle_for_ipc_batched`` -- so nothing reached the client even + though the puts succeeded. + """ + original = RuntimeError("KV cache OOM") + engine = _EngineStub(event_loop_error=original, is_shutdown=True) + helper = _make_helper(engine, num_pending=2, ipc=True) + + assert helper(timeout=0.01) is False + assert helper.handler_kind is AwaitResponseHelper.HandlerKind.ipc_batched + assert not engine._event_loop_error_delivered.is_set() + + def test_gate_stays_clear_when_there_was_nobody_to_wake(self): + """No pending request means nobody is holding the error. + + So the kill must stay armed even in single-process mode. + """ + original = RuntimeError("crash") + engine = _EngineStub(event_loop_error=original, is_shutdown=True) + helper = _make_helper(engine, num_pending=0, ipc=False) + + assert helper(timeout=0.01) is False + assert not engine._event_loop_error_delivered.is_set() + + def test_gate_set_on_the_await_responses_raises_path_too(self): + """The defensive branch also delivers, so it may stand the kill down. + + Previously it never set the gate at all. + """ + original = RuntimeError("unexpected") + engine = _EngineStub(await_responses_raises=original) + helper = _make_helper(engine, num_pending=1, ipc=False) + + assert helper(timeout=0.01) is False + assert engine._event_loop_error_delivered.is_set()