From fcc10dd548ced3356810b9318c1ae1449b4c03fd Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Mon, 3 Aug 2026 13:40:04 +0000 Subject: [PATCH 1/3] [TRTLLM-13409][fix] give the benchmark-disagg fill gate's retry loop a deadline `_check_benchmark_disagg_gate()` retries until the fill completes, with no bound. While it spins the job is invisible: the `continue` it drives sits before `iter_counter += 1`, so the iteration counter freezes while wall-clock advances, and the only outward sign is a stream of byte-identical iteration lines ~110 ms apart -- this gate's own `time.sleep(0.1)` seen from outside. Archived wedges show tens of thousands of them before Slurm kills the job. Bound the retry on LACK OF PROGRESS rather than on elapsed time: a fill that is merely slow keeps resetting the clock and is never killed; only a fill that makes no progress at all for the whole window raises. `TRTLLM_BENCHMARK_DISAGG_FILL_STALL_SEC=0` disables the bound; the default is 600 s. Raising surfaces the stall through the executor loop's existing error path rather than adding a second one, and the rank that raises names itself. Tests (no GPU): the bound fires after the window; the message names the rank and the knob; a slow-but-advancing fill survives 60 s against a 5 s window; progress resets the clock; 0 disables it; and a single no-progress pass only arms the clock rather than raising. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 69 +++++++++++ .../test_disagg_fill_gate_stall_bound.py | 114 ++++++++++++++++++ 2 files changed, 183 insertions(+) create mode 100644 tests/unittest/_torch/executor/test_disagg_fill_gate_stall_bound.py diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 60d801d904c7..63c830dde23a 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -114,6 +114,28 @@ 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: + return 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 + + # 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). @@ -767,6 +789,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 @@ -4050,12 +4076,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. diff --git a/tests/unittest/_torch/executor/test_disagg_fill_gate_stall_bound.py b/tests/unittest/_torch/executor/test_disagg_fill_gate_stall_bound.py new file mode 100644 index 000000000000..946b947ed061 --- /dev/null +++ b/tests/unittest/_torch/executor/test_disagg_fill_gate_stall_bound.py @@ -0,0 +1,114 @@ +# 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 PyExecutor + + +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 From b992d6ab0f9de5528791403bdc7b0ca67b29bb02 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Wed, 12 Aug 2026 06:10:32 +0000 Subject: [PATCH 2/3] [TRTLLM-13409][fix] reject a non-finite fill-gate stall bound, and test on CPU Two follow-ups before review. float() accepts "nan" and "inf", and each breaks the new bound in an opposite direction. Every nan comparison is False, so `nan <= 0` does not disable the bound 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 600s window, which destroys exactly the margin the default exists to give. inf is the mirror: `stalled_for < inf` is always True, so the bound never fires and is silently equivalent to 0. Reject both and fall back to the default, the same guard merged for TLLM_RANK_CRASH_HARD_KILL_GRACE in #16592. The tests are pure monkeypatch with no engine and no GPU, but lacked the cpu_only marker, so the l0_cpu `unittest/_torch/executor` entry collected nothing and they ran only on the h100/b300/gb300 stages. Add the marker so they also run in the CPU stage, where they belong. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 16 ++++++++- .../test_disagg_fill_gate_stall_bound.py | 36 ++++++++++++++++++- 2 files changed, 50 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 63c830dde23a..70998301620a 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/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 @@ -128,12 +129,25 @@ def _fill_stall_timeout_sec() -> float: if raw is None: return _BENCHMARK_DISAGG_FILL_STALL_DEFAULT_SEC try: - return float(raw) + 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. diff --git a/tests/unittest/_torch/executor/test_disagg_fill_gate_stall_bound.py b/tests/unittest/_torch/executor/test_disagg_fill_gate_stall_bound.py index 946b947ed061..2a4fa9cb27f1 100644 --- a/tests/unittest/_torch/executor/test_disagg_fill_gate_stall_bound.py +++ b/tests/unittest/_torch/executor/test_disagg_fill_gate_stall_bound.py @@ -28,7 +28,16 @@ import pytest -from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor +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: @@ -112,3 +121,28 @@ def test_first_stalled_call_only_arms_the_clock(monkeypatch): 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 From 4bf16a657fe968c0237a1ec701b4760e5db23880 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Wed, 12 Aug 2026 11:27:38 +0000 Subject: [PATCH 3/3] [TRTLLM-13409][test] give the fill-gate test double the stall state it now needs The gate now calls _fail_if_fill_gate_stalled(), and MockBenchmarkExecutor binds the gate methods off PyExecutor without having that one or the two attributes it reads. 12 CPU-Generic failures in pipeline 53254, all "AttributeError: 'MockBenchmarkExecutor' object has no attribute '_fail_if_fill_gate_stalled'". Only visible after the rebase: test_benchmark_disagg.py arrived on main after this branch was cut, so the pre-rebase runs never exercised it. Bind the method and seed the state with the bound disabled. That is right twice over: these tests assert retry semantics rather than the deadline, which test_disagg_fill_gate_stall_bound.py covers, and they patch the whole `time` module -- an enabled bound would do arithmetic on a Mock. The `timeout_s <= 0` guard returns before the clock is read, so the patched module is never touched. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- .../unittest/_torch/executor/test_benchmark_disagg.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/unittest/_torch/executor/test_benchmark_disagg.py b/tests/unittest/_torch/executor/test_benchmark_disagg.py index e800992282ab..71cba9b88734 100644 --- a/tests/unittest/_torch/executor/test_benchmark_disagg.py +++ b/tests/unittest/_torch/executor/test_benchmark_disagg.py @@ -104,6 +104,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) @@ -112,6 +121,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 # ---------------------------------------------------------------------------