[TRTLLM-13409][fix] give the benchmark-disagg fill gate's retry loop a deadline

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -4,6 +4,7 @@
 import dataclasses
 import datetime
 import functools
+import math
 import os
 import sys
 import threading
@@ -124,6 +125,41 @@ def _stats_buffer_is_unbounded(max_stats_len: int) -> bool:
 # Environment variable to control the benchmark disagg fill target.
 BENCHMARK_REQ_QUEUES_SIZE_ENV_VAR_NAME = "TLLM_BENCHMARK_REQ_QUEUES_SIZE"
 
+# How long the benchmark-disagg fill gate may retry without making ANY
+# transfer progress before the rank gives up. Generous by default: a fill that
+# is merely slow keeps resetting the clock, so this only fires on a fill that
+# is not advancing at all. 0 disables the bound.
+BENCHMARK_DISAGG_FILL_STALL_ENV_VAR_NAME = "TRTLLM_BENCHMARK_DISAGG_FILL_STALL_SEC"
+_BENCHMARK_DISAGG_FILL_STALL_DEFAULT_SEC = 600.0
+
+
+def _fill_stall_timeout_sec() -> float:
+    """Resolve the fill-gate stall bound; malformed values fall back."""
+    raw = os.environ.get(BENCHMARK_DISAGG_FILL_STALL_ENV_VAR_NAME)
+    if raw is None:
+        return _BENCHMARK_DISAGG_FILL_STALL_DEFAULT_SEC
+    try:
+        value = float(raw)
+    except ValueError:
+        logger.warning(
+            f"Invalid {BENCHMARK_DISAGG_FILL_STALL_ENV_VAR_NAME}={raw!r}; "
+            f"using {_BENCHMARK_DISAGG_FILL_STALL_DEFAULT_SEC:.0f}s")
+        return _BENCHMARK_DISAGG_FILL_STALL_DEFAULT_SEC
+    if not math.isfinite(value):
+        # float() accepts "nan" and "inf", and both break the bound in
+        # opposite directions. Every nan comparison is False, so `nan <= 0`
+        # does not disable it and `stalled_for < nan` does not defer it --
+        # control falls straight through to the raise, firing on the second
+        # consecutive stalled call (~0.1s) instead of after the window. inf
+        # is the mirror: `stalled_for < inf` is always True, so the bound
+        # never fires and is silently equivalent to 0. Reject both.
+        logger.warning(
+            f"Non-finite {BENCHMARK_DISAGG_FILL_STALL_ENV_VAR_NAME}={raw!r}; "
+            f"using {_BENCHMARK_DISAGG_FILL_STALL_DEFAULT_SEC:.0f}s")
+        return _BENCHMARK_DISAGG_FILL_STALL_DEFAULT_SEC
+    return value
+
+
 # Environment variable to control which ranks print step logging.
 # Format: comma-separated rank IDs, e.g. "0,1,3", or "all" for all ranks.
 # Default: "0" (only rank 0 prints, matching existing behavior).
@@ -848,6 +884,10 @@ def __init__(
         self.has_previous_draft_tokens = False
         self.num_scheduled_requests: int = 0
         self._configure_benchmark_req_queues_size()
+        # Deadline state for the benchmark-disagg fill gate's retry loop.
+        # None means "not currently stalled"; progress resets it.
+        self._benchmark_fill_stall_since: Optional[float] = None
+        self._benchmark_fill_stall_timeout_sec = _fill_stall_timeout_sec()
 
         # Sample-state relay mode. "1": a background thread relays them,
         # overlapping with forward, but it needs the GIL and can be starved
@@ -4194,12 +4234,55 @@ def _check_benchmark_disagg_gate(self, scheduled_batch: ScheduledRequests,
             if can_forward:
                 self._benchmark_fill_phase_active = False
                 self._fill_admit_cap = 0
+                self._benchmark_fill_stall_since = None
             elif not sync_transfer_made_progress:
                 time.sleep(0.1)
             if not can_forward:
+                self._fail_if_fill_gate_stalled(sync_transfer_made_progress)
                 return can_forward, True
         return can_forward, False
 
+    def _fail_if_fill_gate_stalled(self, made_progress: bool) -> None:
+        """Give the fill gate's retry loop a deadline.
+
+        The retry is otherwise unbounded, and it is invisible while it spins:
+        the ``continue`` it drives sits before ``iter_counter += 1``, so the
+        iteration counter freezes while wall-clock advances. Archived wedges
+        show exactly that -- tens of thousands of byte-identical iteration
+        lines at ``host_step_time`` ~110 ms, which is this function's own
+        ``time.sleep(0.1)`` seen from outside.
+
+        Progress resets the clock, so a slow-but-advancing fill is never
+        killed; only a fill that makes no progress at all for the whole
+        window is. ``TRTLLM_BENCHMARK_DISAGG_FILL_STALL_SEC=0`` disables the
+        bound entirely.
+
+        Raising here surfaces the stall through the executor loop's normal
+        error path rather than adding a second one, and the rank that raises
+        names itself in the message.
+        """
+        if made_progress:
+            self._benchmark_fill_stall_since = None
+            return
+        timeout_s = self._benchmark_fill_stall_timeout_sec
+        if timeout_s <= 0:
+            return
+        now = time.monotonic()
+        if self._benchmark_fill_stall_since is None:
+            self._benchmark_fill_stall_since = now
+            return
+        stalled_for = now - self._benchmark_fill_stall_since
+        if stalled_for < timeout_s:
+            return
+        self._benchmark_fill_stall_since = None
+        raise RuntimeError(
+            f"Benchmark disagg fill gate made no progress for "
+            f"{stalled_for:.0f}s on rank {self.dist.rank} "
+            f"(TRTLLM_BENCHMARK_DISAGG_FILL_STALL_SEC="
+            f"{timeout_s:.0f}). The KV transfers this gate waits on are not "
+            f"completing; the loop would otherwise retry forever without "
+            f"advancing iter_counter.")
+
     @nvtx_range("_handle_disagg_cache_errors_synced")
     def _handle_disagg_cache_errors_synced(self):
         """Rank-safe disagg cache error and poison handler.

tests/unittest/_torch/executor/test_benchmark_disagg.py

@@ -113,6 +113,15 @@ def __init__(
         self.dist.tp_size = tp_size
         self.dist.world_size = tp_size
 
+        # State the gate's stall bound reads. 0 disables the bound, which is
+        # what these tests want twice over: they assert retry semantics, not
+        # the deadline (that is test_disagg_fill_gate_stall_bound.py's job),
+        # and they patch the whole `time` module -- so an enabled bound would
+        # do arithmetic on a Mock. `timeout_s <= 0` returns before the clock
+        # is read, so the patched module is never touched.
+        self._benchmark_fill_stall_since = None
+        self._benchmark_fill_stall_timeout_sec = 0.0
+
     from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor
 
     _dist_size = staticmethod(PyExecutor._dist_size)
@@ -121,6 +130,7 @@ def __init__(
     _configure_benchmark_req_queues_size = PyExecutor._configure_benchmark_req_queues_size
     _is_benchmark_disagg_fill_complete = PyExecutor._is_benchmark_disagg_fill_complete
     _check_benchmark_disagg_gate = PyExecutor._check_benchmark_disagg_gate
+    _fail_if_fill_gate_stalled = PyExecutor._fail_if_fill_gate_stalled
 
 
 # ---------------------------------------------------------------------------

tests/unittest/_torch/executor/test_disagg_fill_gate_stall_bound.py

@@ -0,0 +1,148 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""The benchmark-disagg fill gate's retry loop must have a deadline.
+
+The gate retries until the fill completes. Its `continue` sits before
+`iter_counter += 1`, so while it spins the iteration counter is frozen and
+wall-clock advances -- which is why an archived wedge shows tens of thousands
+of byte-identical iteration lines at ~110 ms apart (this gate's own
+`time.sleep(0.1)`, observed from outside).
+
+The bound is on *no progress*, not on elapsed time: a slow but advancing fill
+resets the clock and is never killed.
+"""
+
+import types
+
+import pytest
+
+from tensorrt_llm._torch.pyexecutor.py_executor import (
+    BENCHMARK_DISAGG_FILL_STALL_ENV_VAR_NAME,
+    PyExecutor,
+    _fill_stall_timeout_sec,
+)
+
+# Pure monkeypatch: no engine, no GPU. The marker is required, not
+# decorative -- tests/unittest/conftest.py's pytest_ignore_collect drops any
+# file whose source lacks this literal when pytest runs with -m cpu_only.
+pytestmark = pytest.mark.cpu_only
+
+
+class _Clock:
+    def __init__(self):
+        self.t = 1000.0
+
+    def __call__(self):
+        return self.t
+
+    def advance(self, dt):
+        self.t += dt
+
+
+def _executor(timeout_s, clock, rank=0):
+    """A PyExecutor with only the fill-gate stall surface populated."""
+    ex = object.__new__(PyExecutor)
+    ex._benchmark_fill_stall_since = None
+    ex._benchmark_fill_stall_timeout_sec = timeout_s
+    ex.dist = types.SimpleNamespace(rank=rank)
+    return ex
+
+
+def _spin(ex, clock, monkeypatch, seconds, step=0.1, made_progress=False):
+    """Drive the no-progress retry path for `seconds` of wall clock."""
+    monkeypatch.setattr("tensorrt_llm._torch.pyexecutor.py_executor.time.monotonic", clock)
+    for _ in range(int(seconds / step)):
+        ex._fail_if_fill_gate_stalled(made_progress)
+        clock.advance(step)
+
+
+def test_raises_once_the_stall_window_elapses(monkeypatch):
+    clock = _Clock()
+    ex = _executor(5.0, clock, rank=3)
+    with pytest.raises(RuntimeError, match="made no progress"):
+        _spin(ex, clock, monkeypatch, seconds=8.0)
+
+
+def test_message_names_the_rank_and_the_knob(monkeypatch):
+    clock = _Clock()
+    ex = _executor(5.0, clock, rank=7)
+    with pytest.raises(RuntimeError) as exc:
+        _spin(ex, clock, monkeypatch, seconds=8.0)
+    msg = str(exc.value)
+    assert "rank 7" in msg
+    assert "TRTLLM_BENCHMARK_DISAGG_FILL_STALL_SEC" in msg
+
+
+def test_a_slow_but_advancing_fill_is_never_killed(monkeypatch):
+    """The bound is on no-progress, not on elapsed time.
+
+    A fill that keeps making progress may legitimately take far longer than
+    the window; killing it would be a regression, not a fix.
+    """
+    clock = _Clock()
+    ex = _executor(5.0, clock)
+    _spin(ex, clock, monkeypatch, seconds=60.0, made_progress=True)
+    assert ex._benchmark_fill_stall_since is None
+
+
+def test_progress_resets_the_clock(monkeypatch):
+    """Stall, recover, stall again -- the second window starts from zero."""
+    clock = _Clock()
+    ex = _executor(5.0, clock)
+    _spin(ex, clock, monkeypatch, seconds=4.0)  # just under
+    ex._fail_if_fill_gate_stalled(True)  # progress
+    assert ex._benchmark_fill_stall_since is None
+    _spin(ex, clock, monkeypatch, seconds=4.0)  # under again
+    assert ex._benchmark_fill_stall_since is not None  # armed, not fired
+
+
+def test_zero_disables_the_bound(monkeypatch):
+    clock = _Clock()
+    ex = _executor(0.0, clock)
+    _spin(ex, clock, monkeypatch, seconds=3600.0, step=10.0)
+    assert ex._benchmark_fill_stall_since is None
+
+
+def test_first_stalled_call_only_arms_the_clock(monkeypatch):
+    """One no-progress pass is normal; it must not raise on its own."""
+    monkeypatch.setattr("tensorrt_llm._torch.pyexecutor.py_executor.time.monotonic", _Clock())
+    ex = _executor(5.0, _Clock())
+    ex._fail_if_fill_gate_stalled(False)
+    assert ex._benchmark_fill_stall_since is not None
+
+
+@pytest.mark.parametrize("raw", ["nan", "NaN", "inf", "-inf", "Infinity"])
+def test_non_finite_env_falls_back_to_the_default(monkeypatch, raw):
+    """float() accepts nan/inf, and each breaks the bound a different way.
+
+    Every nan comparison is False, so ``nan <= 0`` does not disable the
+    bound and ``stalled_for < nan`` does not defer it: control falls through
+    to the raise, firing on the second consecutive stalled call rather than
+    after the window. ``inf`` is the mirror -- ``stalled_for < inf`` is
+    always True, so the bound never fires and is silently equivalent to 0.
+    """
+    monkeypatch.setenv(BENCHMARK_DISAGG_FILL_STALL_ENV_VAR_NAME, raw)
+    assert _fill_stall_timeout_sec() == 600.0
+
+
+def test_a_valid_env_override_is_honoured(monkeypatch):
+    monkeypatch.setenv(BENCHMARK_DISAGG_FILL_STALL_ENV_VAR_NAME, "42.5")
+    assert _fill_stall_timeout_sec() == 42.5
+
+
+def test_zero_env_is_preserved_as_the_disable_switch(monkeypatch):
+    """0 must survive the finiteness check -- it is the documented opt-out."""
+    monkeypatch.setenv(BENCHMARK_DISAGG_FILL_STALL_ENV_VAR_NAME, "0")
+    assert _fill_stall_timeout_sec() == 0.0