[TRTLLM-13409][fix] hard-kill all ranks when one rank's executor loop crashes

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,190 @@ 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,
+) -> 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
+        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):
+        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()
+
+    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)
+
+
+def start_rank_crash_kill_watchdog(world_size: int) -> 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)
+        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.
 

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,
@@ -822,6 +823,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()
@@ -1185,6 +1195,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
@@ -1194,6 +1206,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
@@ -1206,7 +1232,36 @@ def _event_loop_wrapper(self):
             self._event_loop_error = e
             raise e
         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) 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)
 
     @property
     def is_warmup(self) -> bool:
@@ -2546,6 +2601,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()
@@ -2831,6 +2887,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()
@@ -4012,6 +4074,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(
@@ -4488,6 +4551,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(