From 80f6718dde6d1642d8627bba94a02929c0ba0a75 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Tue, 11 Aug 2026 12:13:55 -0700 Subject: [PATCH 1/3] [None][perf] Arm the hang detector once instead of per checkpoint checkpoint() cancelled the pending watchdog task and scheduled a new one, and both calls post a call_soon_threadsafe, so the executor loop woke the detector thread several times per iteration to re-arm a 300s watchdog. Keep one long-lived watcher and move a monotonic deadline instead. The timeout is now measured from the checkpoint call rather than from when the detector thread gets to process it. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- .../_torch/pyexecutor/hang_detector.py | 69 ++++++++++-- .../executor/test_hang_detector_kill.py | 103 +++++++++++++++++- 2 files changed, 161 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/hang_detector.py b/tensorrt_llm/_torch/pyexecutor/hang_detector.py index 9773ff324849..9fb7e56ab59d 100644 --- a/tensorrt_llm/_torch/pyexecutor/hang_detector.py +++ b/tensorrt_llm/_torch/pyexecutor/hang_detector.py @@ -346,6 +346,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 +357,67 @@ 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 + self.active = True + + # Disarmed until the first checkpoint so startup does not lapse. + self._deadline = math.inf 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 never cancels or + # re-arms it; it 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. While disarmed + the deadline is ``inf``; the 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 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. + """ + while self.active: + deadline = self._deadline + remaining = deadline - time.monotonic() + if remaining > 0: + await asyncio.sleep(min(remaining, self.timeout)) + continue + # Disarm only the deadline observed to lapse, so one lapse reports + # once. A checkpoint racing this branch installs a newer deadline; + # clearing that would leave the watcher alive but permanently + # disarmed, which is the failure this watchdog exists to catch. + if self._deadline == deadline: + self._deadline = math.inf + 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 +451,18 @@ 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 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 + """Disarm hang detection until the next checkpoint.""" + self._deadline = math.inf @contextmanager def pause(self): """Pause hang detection in scope.""" + self._deadline = math.inf try: - self.cancel_task() yield finally: self.checkpoint() diff --git a/tests/unittest/_torch/executor/test_hang_detector_kill.py b/tests/unittest/_torch/executor/test_hang_detector_kill.py index 210c3bbbda82..8d4f003b04cf 100644 --- a/tests/unittest/_torch/executor/test_hang_detector_kill.py +++ b/tests/unittest/_torch/executor/test_hang_detector_kill.py @@ -16,6 +16,7 @@ import asyncio import contextlib +import math import os import shutil import signal @@ -66,6 +67,106 @@ def test_checkpoint_resets_timer(): assert hd.detected() is False +def test_checkpoint_reuses_one_watcher_task(): + """One watcher task serves every checkpoint, pause and resume. + + The executor loop checkpoints several times per iteration, and each + schedule/cancel of a task wakes the detector's event-loop thread, so the + single-task design is what keeps checkpoint() off that thread entirely. + """ + hd = HangDetector(timeout=30) + with hd: + task = hd.task + assert task is not None + for _ in range(10): + hd.checkpoint() + with hd.pause(): + hd.checkpoint() + hd.checkpoint() + assert hd.task is task + assert not task.done() + + +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. + assert not hd.task.done() + 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_a_checkpoint_racing_the_disarm_is_not_erased(monkeypatch): + """A checkpoint landing as the watcher disarms must survive. + + The watcher reads the lapsed deadline, then clears it. A checkpoint in + between installs a newer deadline; clearing that would leave the watcher + running but permanently disarmed, so the work it just armed could hang + undetected. + """ + hd = HangDetector(timeout=1, on_detected=lambda: None) + real_monotonic = hang_detector_module.time.monotonic + state = {"injected": False, "busy": False} + + def racing_monotonic(): + now = real_monotonic() + if state["injected"] or state["busy"]: + return now + # Only inject from `_watch`'s own lapse computation. asyncio's event + # loop also reads the clock, and injecting from there would land + # outside the window and silently make this test vacuous. + caller = sys._getframe(1) + if caller.f_code.co_name != "_watch" or hd._deadline > now: + return now + # `_watch` has already read `self._deadline` into a local by now, so + # this checkpoint lands exactly between that read and the disarm. + state["busy"] = True + state["injected"] = True + hd.checkpoint() + state["busy"] = False + return now + + monkeypatch.setattr(hang_detector_module.time, "monotonic", racing_monotonic) + with hd: + hd.checkpoint() + deadline = real_monotonic() + 5.0 + while real_monotonic() < deadline and not state["injected"]: + time.sleep(0.05) + assert state["injected"], "the racing checkpoint never landed" + time.sleep(0.2) # let the watcher finish its disarm/report pass + assert hd._deadline != math.inf, "the racing checkpoint's arm was erased" + + def test_pause_suppresses_detection(): fired = [] hd = HangDetector(timeout=1, on_detected=lambda: fired.append(1)) @@ -102,7 +203,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 From 9d80c94b09f7764955acac2df3419e756a4a24bd Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Tue, 11 Aug 2026 12:14:07 -0700 Subject: [PATCH 2/3] [None][perf] Skip recording profiler timing events nobody reads The start events were recorded on every executor iteration while the read side is already gated on print_log or enable_iter_perf_stats, so with neither enabled the record calls were pure overhead. Gate the record on the same condition. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 00037596ad4d..87295bd20c59 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -1813,9 +1813,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() @@ -1880,14 +1879,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 From ab284c71e621e7dfcdd44db35c4b5bab82a89477 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Wed, 12 Aug 2026 11:11:56 -0700 Subject: [PATCH 3/3] [None][fix] Disarm the hang detector before publishing active start() published `active` inside the lock but reset `_deadline` after releasing it. A checkpoint from the executor thread in that window armed the deadline, and the reset then erased it, leaving the loop unwatched until the next checkpoint. Store the disarmed deadline inside the lock, before `active` goes true. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/hang_detector.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/hang_detector.py b/tensorrt_llm/_torch/pyexecutor/hang_detector.py index 9fb7e56ab59d..bb191a6f5932 100644 --- a/tensorrt_llm/_torch/pyexecutor/hang_detector.py +++ b/tensorrt_llm/_torch/pyexecutor/hang_detector.py @@ -367,10 +367,12 @@ def run_loop(): "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 - # Disarmed until the first checkpoint so startup does not lapse. - self._deadline = math.inf self.loop = asyncio.new_event_loop() self.loop_thread = threading.Thread(target=run_loop, daemon=True, name="hang_detector_loop") self.loop_thread.start()