[None][perf] Cut per-iteration executor bookkeeping in the hang detector and profiler

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,29 +478,36 @@ 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()
 
     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():

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

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