[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."""
@@ -80,6 +86,100 @@ def propagate_hard_kill(exit_code: int = _HARD_KILL_EXIT_CODE) -> None:
     os.kill(os.getpid(), signal.SIGKILL)
 
 
+def _rank_crash_kill_grace() -> Optional[float]:
+    """Resolve the crash-kill grace period; ``None`` means the kill is disabled."""
+    raw = os.environ.get(RANK_CRASH_KILL_GRACE_ENV)
+    if raw is None:
+        return _RANK_CRASH_KILL_GRACE_DEFAULT
+    try:
+        grace = float(raw)
+    except ValueError:
+        _best_effort_log_error(
+            f"Invalid {RANK_CRASH_KILL_GRACE_ENV}={raw!r}; "
+            f"using default {_RANK_CRASH_KILL_GRACE_DEFAULT}s"
+        )
+        return _RANK_CRASH_KILL_GRACE_DEFAULT
+    return None if grace < 0 else grace
+
+
+def hard_kill_on_rank_crash(world_size: int) -> bool:
+    """Hard-kill the whole world after this rank's executor loop crashed.
+
+    A rank whose executor loop died on an exception can never rejoin its
+    peers' collectives: without an explicit kill, every peer blocks in its
+    next collective until its own HangDetector fires (300 s), and the whole
+    test session burns that long for an error that was already known.
+
+    The grace sleep before the kill is load-bearing: it gives the crashed
+    rank's cleaner error paths time to win the race, so the client reports
+    the ORIGINAL exception instead of a bare worker death —
+    - rank-local response waiters woken by the executor-loop cleanup read
+      the stashed error and surface it through the response path;
+    - during init, the worker's ready handshake returns the real error to
+      the proxy before the abort tears the world down;
+    - the worker main thread returning lets its mpi4py future complete with
+      the original exception.
+
+    Never raises (it runs in a ``finally`` where an exception would mask the
+    original loop error). Returns True when the kill path was taken — only
+    observable in tests, where ``propagate_hard_kill`` is stubbed; in
+    production that call does not return.
+    """
+    try:
+        if world_size <= 1:
+            # No peers to unblock; the worker's own death already completes
+            # its future/handshake with the original exception.
+            return False
+        grace = _rank_crash_kill_grace()
+        if grace is None:
+            return False
+        _best_effort_log_error(
+            f"Executor loop crashed on this rank; hard-killing all "
+            f"{world_size} ranks in {grace}s (peers cannot make progress "
+            f"without this rank). Set {RANK_CRASH_KILL_GRACE_ENV}=-1 to disable."
+        )
+        if grace > 0:
+            time.sleep(grace)
+        propagate_hard_kill()
+        return True
+    except Exception as e:  # noqa: BLE001 - must not mask the loop's original error
+        _best_effort_log_error(f"hard_kill_on_rank_crash failed (ignored): {e!r}")
+        return False
+
+
+def start_rank_crash_kill_watchdog(world_size: int) -> Optional[threading.Thread]:
+    """Arm a daemon thread that hard-kills the world once the grace elapses.
+
+    Must be armed BEFORE executor-loop cleanup: cleanup can block without
+    bound (e.g. ``wait()`` on a pending PP send handle wedged by the crash),
+    and a kill placed after it would never be reached — leaving peers to
+    burn in their own 300 s HangDetectors, the exact failure this kill
+    exists to avoid. The thread reuses ``hard_kill_on_rank_crash``, so the
+    kill fires at crash + grace whether cleanup finishes, blocks, or raises.
+
+    Never raises. Returns the armed thread, or ``None`` when the kill is
+    not applicable (single rank, disabled by env) or the thread could not
+    be started — in that case the caller's post-cleanup kill remains the
+    only mechanism.
+    """
+    try:
+        if world_size <= 1:
+            return None
+        if _rank_crash_kill_grace() is None:
+            return None
+        watchdog = threading.Thread(
+            target=hard_kill_on_rank_crash,
+            args=(world_size,),
+            name="rank_crash_kill_watchdog",
+            daemon=True,
+        )
+        watchdog.start()
+        return watchdog
+    except Exception as e:  # noqa: BLE001 - must not mask the loop's original error
+        _best_effort_log_error(f"failed to arm rank-crash kill watchdog (ignored): {e!r}")
+        return None
+
+
 class HangDetector:
     """Watchdog that fires when the executor loop stops checkpointing.
 

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -66,7 +66,8 @@
 from .guided_decoder import GuidedDecoder
 from .handle_additional_outputs import HandleAdditionalOutputs
 from .handle_logits import HandleLogits
-from .hang_detector import HangDetector, propagate_hard_kill
+from .hang_detector import (HangDetector, hard_kill_on_rank_crash,
+                            propagate_hard_kill, start_rank_crash_kill_watchdog)
 from .kv_cache_manager_v2 import KVCacheManagerV2
 from .kv_cache_stats import append_kv_cache_iteration_stats
 from .kv_cache_transceiver import (KvCacheTransceiver,
@@ -1180,6 +1181,7 @@ def _flush_iter_stats_synced(self):
     # Performance metrics methods are in PerfMetricsManager (self.perf_manager)
 
     def _event_loop_wrapper(self):
+        crashed = False
         try:
             # Skip line profiler during warmup/memory estimation phase to avoid
             # saving incomplete results that would be overwritten anyway
@@ -1189,6 +1191,7 @@ def _event_loop_wrapper(self):
                  customized_gc_thresholds(self.garbage_collection_gen0_threshold):
                 self.event_loop()
         except Exception as e:
+            crashed = True
             logger.error(f"Error in event loop: {e}")
             logger.error(traceback.format_exc())
             # Stash the original error so local consumers
@@ -1201,7 +1204,25 @@ def _event_loop_wrapper(self):
             self._event_loop_error = e
             raise e
         finally:
-            self._executor_loop_cleanup()
+            if crashed:
+                # Armed BEFORE cleanup: cleanup can block without bound on a
+                # send handle wedged by the crash, and a kill placed only
+                # after it would never fire.
+                start_rank_crash_kill_watchdog(self.dist.world_size)
+            try:
+                self._executor_loop_cleanup()
+            finally:
+                if crashed:
+                    # Peers cannot make progress without this rank's loop:
+                    # they would block in their next collective until their
+                    # own HangDetectors fire 300s later. Kill the world now
+                    # instead; the grace inside lets the stashed error reach
+                    # rank-local waiters and the ready handshake first, so
+                    # the client sees the original exception rather than a
+                    # bare worker death. Nested finally: the kill must fire
+                    # even when cleanup itself raises; the watchdog above
+                    # covers cleanup blocking.
+                    hard_kill_on_rank_crash(self.dist.world_size)
 
     @property
     def is_warmup(self) -> bool: