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

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.

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.