diff --git a/tensorrt_llm/_torch/pyexecutor/hang_detector.py b/tensorrt_llm/_torch/pyexecutor/hang_detector.py index 9773ff324849..06dcdbca3b1e 100644 --- a/tensorrt_llm/_torch/pyexecutor/hang_detector.py +++ b/tensorrt_llm/_torch/pyexecutor/hang_detector.py @@ -328,9 +328,23 @@ def start_rank_crash_kill_watchdog( class HangDetector: """Watchdog that fires when the executor loop stops checkpointing. - When ``timeout`` seconds pass without a ``checkpoint()``, all thread stacks - are dumped for diagnosis and ``on_detected`` runs (the hard-kill + - cross-rank propagation path). + Contract: + + - ``timeout`` seconds without a ``checkpoint()`` dumps all thread stacks for + diagnosis and runs ``on_detected`` (the hard-kill + cross-rank + propagation path). + - Continued checkpointing never fires it. A false positive hard-kills a + healthy job, so this bound is as load-bearing as detection itself. + - ``start()`` leaves detection disarmed; the first ``checkpoint()`` arms it, + so the start-to-first-checkpoint window is not hang-eligible. + - ``pause()`` suppresses detection in scope and re-arms on exit. It does + not nest: leaving an inner ``pause()`` re-arms while an outer one is + still open. + - Detection never stops while active: not after firing, and not if + ``on_detected`` raises an ``Exception``. ``on_detected`` is not + idempotent, so a single lapse invokes it once. + - ``checkpoint()`` is one clock read and one float store, and does no + cross-thread work. The executor loop calls it three times per iteration. """ def __init__( @@ -346,6 +360,9 @@ def __init__( self.active = False self._detected = False self._status_providers: list[Callable[[], str]] = [] + # Monotonic stamp the watcher compares against; ``inf`` means disarmed. + # A plain float store is the entire cost of ``checkpoint()``. + self._deadline = math.inf def start(self): """Enable hang detection.""" @@ -354,18 +371,80 @@ def run_loop(): asyncio.set_event_loop(self.loop) self.loop.run_forever() - self.active = True + with self.lock: + # Locked, not a bare check: concurrent callers could both observe + # ``active`` false and schedule a watcher, and watchers share + # ``_deadline``, so a second one reports the same lapse twice and + # propagates two hard kills. + if self.active: + _best_effort_log_error( + "HangDetector.start() called while already active; ignoring." + ) + return + # Disarmed until the first checkpoint so startup does not lapse. + # Stored before ``active`` is published so a checkpoint racing this + # call cannot have its arm overwritten here. + self._deadline = math.inf + self.active = True + self.loop = asyncio.new_event_loop() self.loop_thread = threading.Thread(target=run_loop, daemon=True, name="hang_detector_loop") self.loop_thread.start() + # One long-lived watcher, scheduled once; the hot path only moves + # ``_deadline``. + self.task = asyncio.run_coroutine_threadsafe(self._watch(), self.loop) def register_status_provider(self, provider: Callable[[], str]) -> None: """Register a nonblocking callable that returns status to dump on hang detection.""" with self.lock: self._status_providers.append(provider) - async def _detect_hang(self) -> None: - await asyncio.sleep(self.timeout) + async def _watch(self) -> None: + """Sleep until the deadline lapses, report, and keep watching. + + Waking early is normal: ``checkpoint()`` pushes ``_deadline`` forward + without touching this task, so each wake-up either finds time left and + sleeps again, or finds the deadline passed and reports. Every sleep is + clamped to ``timeout`` because ``checkpoint()`` only stores a float and + never wakes this loop, so an unclamped sleep would not notice a later + arm. + + This task never writes ``_deadline``; the lapse it last reported is + watcher-local. + + This task outlives a report, and outlives a report that raises. A + watchdog that quietly stopped watching would be the exact failure it + exists to catch, and ``on_detected`` is the cross-rank hard kill, which + can itself fail on an already-degraded job. + """ + # The deadline whose lapse already ran ``on_detected``. Compared by + # identity, not equality: ``checkpoint()`` publishes a fresh float, so a + # re-arm that lands on the same value still reports. + reported = None + while self.active: + # Clock first: a checkpoint can land between these two reads, and + # whichever is read first is the stale one. A stale deadline fires + # at work that was checkpointed in time; a stale clock only defers. + now = time.monotonic() + deadline = self._deadline + remaining = deadline - now + if remaining > 0: + await asyncio.sleep(min(remaining, self.timeout)) + continue + if deadline is reported: + # ``on_detected`` is not idempotent, so one lapse runs it once. + # Nothing wakes this loop, so poll for the next arm. + await asyncio.sleep(self.timeout) + continue + reported = deadline + try: + await self._report_hang() + except Exception as error: # noqa: BLE001 - the watcher must survive + _best_effort_log_error( + f"HangDetector: reporting failed with {type(error).__name__}: {error}" + ) + + async def _report_hang(self) -> None: with self.lock: status_providers = tuple(self._status_providers) @@ -399,21 +478,28 @@ def detected(self): def checkpoint(self): """Reset hang detection timer.""" - self.cancel_task() if self.active: - self.task = asyncio.run_coroutine_threadsafe(self._detect_hang(), self.loop) + self._deadline = time.monotonic() + self.timeout + + def disarm(self) -> None: + """Disarm hang detection until the next checkpoint.""" + self._deadline = math.inf + + def cancel_task(self) -> None: + """Compatibility alias for :meth:`disarm`. - def cancel_task(self): - """Cancel the hang detection task.""" - if self.task is not None and not self.task.done(): - self.task.cancel() - self.task = None + The watcher is long-lived and has no task to cancel, but the old name is + load-bearing for the cache-transceiver precheck and its SLURM example. + Delegating rather than aliasing keeps a subclass override of ``disarm`` + effective through this name. + """ + self.disarm() @contextmanager def pause(self): """Pause hang detection in scope.""" + self.disarm() try: - self.cancel_task() yield finally: self.checkpoint() @@ -421,7 +507,7 @@ def pause(self): def stop(self): """Stop hang detection.""" self.active = False - self.cancel_task() + self.disarm() if self.loop is not None: # Cancel all pending tasks before stopping the loop def cancel_all_tasks(): diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 7838491f1e8a..5a1369923c42 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -1903,9 +1903,8 @@ def profile_step(): # — the events being read have already passed by the time we # read them. Stashing on self lets the /metrics serializer pick # up the values without going through the log line. - should_capture_timing = start_time is not None and ( - self.print_log or self.enable_iter_perf_stats) - if should_capture_timing: + should_capture_timing = self.print_log or self.enable_iter_perf_stats + if should_capture_timing and start_time is not None: end_time = time.time() if it % 2 == 0: end_event_1.record() @@ -1970,14 +1969,15 @@ def profile_step(): calibrator.pre_step(it) start_time = time.time() - if it % 2 == 0: - if start_event_1 is None: - start_event_1 = torch.cuda.Event(enable_timing=True) - start_event_1.record() - else: - if start_event_2 is None: - start_event_2 = torch.cuda.Event(enable_timing=True) - start_event_2.record() + if should_capture_timing: + if it % 2 == 0: + if start_event_1 is None: + start_event_1 = torch.cuda.Event(enable_timing=True) + start_event_1.record() + else: + if start_event_2 is None: + start_event_2 = torch.cuda.Event(enable_timing=True) + start_event_2.record() try: yield profile_step diff --git a/tests/unittest/_torch/executor/test_hang_detector_kill.py b/tests/unittest/_torch/executor/test_hang_detector_kill.py index 210c3bbbda82..463a2e9d0963 100644 --- a/tests/unittest/_torch/executor/test_hang_detector_kill.py +++ b/tests/unittest/_torch/executor/test_hang_detector_kill.py @@ -66,6 +66,93 @@ def test_checkpoint_resets_timer(): assert hd.detected() is False +def test_checkpoint_schedules_no_work_on_the_detector_loop( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """checkpoint() runs on the executor thread and must stay off the detector loop. + + Waking that loop is the per-iteration cost this detector is built to avoid, + and the executor pays it three times per iteration. Both routes into the + loop -- scheduling a coroutine and cancelling one -- funnel through + call_soon_threadsafe, so counting it catches either. + """ + hd = HangDetector(timeout=30) + with hd: + woken = [] + real_call_soon_threadsafe = hd.loop.call_soon_threadsafe + + def counting_call_soon_threadsafe(*args: object, **kwargs: object) -> asyncio.Handle: + woken.append(args[0] if args else None) + return real_call_soon_threadsafe(*args, **kwargs) + + monkeypatch.setattr(hd.loop, "call_soon_threadsafe", counting_call_soon_threadsafe) + + for _ in range(10): + hd.checkpoint() + with hd.pause(): + hd.checkpoint() + hd.checkpoint() + hd.disarm() + assert woken == [] + + +def test_detector_is_disarmed_until_the_first_checkpoint(): + """start() enables detection; the first checkpoint arms the deadline. + + Callers separate lifecycle start from arming, so the start-to-first- + checkpoint window must not be attributed to the loop as a hang. + """ + fired = [] + hd = HangDetector(timeout=1, on_detected=lambda: fired.append(1)) + with hd: + time.sleep(2.0) # would fire if start() armed the deadline itself + assert fired == [] + assert hd.detected() is False + + +def test_watcher_survives_a_raising_callback(): + """on_detected is the cross-rank hard kill and can fail on a broken job.""" + fired = [] + + def boom(): + fired.append(1) + raise RuntimeError("hard kill failed") + + hd = HangDetector(timeout=1, on_detected=boom) + with hd: + hd.checkpoint() + deadline = time.monotonic() + 5.0 + while time.monotonic() < deadline and len(fired) < 1: + time.sleep(0.05) + assert len(fired) == 1 + + # The watcher is still live and still able to report. + hd.checkpoint() + deadline = time.monotonic() + 5.0 + while time.monotonic() < deadline and len(fired) < 2: + time.sleep(0.05) + assert len(fired) == 2 + + +def test_one_lapse_invokes_on_detected_once() -> None: + """A single lapse must not re-run on_detected as the watcher keeps polling. + + Nothing clears the deadline once its lapse is reported, so every later poll + observes the same lapse. on_detected is propagate_hard_kill(), which is not + idempotent -- re-running it would re-propagate to peer ranks. + """ + fired = [] + hd = HangDetector(timeout=1, on_detected=lambda: fired.append(1)) + with hd: + hd.checkpoint() + deadline = time.monotonic() + 5.0 + while time.monotonic() < deadline and not fired: + time.sleep(0.05) + assert len(fired) == 1 + time.sleep(2.5) # several further polls over the same lapse + assert len(fired) == 1, "the same lapse ran on_detected again" + + def test_pause_suppresses_detection(): fired = [] hd = HangDetector(timeout=1, on_detected=lambda: fired.append(1)) @@ -102,7 +189,7 @@ def failing_provider(): detector.register_status_provider(failing_provider) detector.register_status_provider(lambda: "transceiver status") - asyncio.run(detector._detect_hang()) + asyncio.run(detector._report_hang()) messages = "\n".join(message for kind, message in events if kind == "log") assert "provider failed" in messages