From 76317101144eac67dc2737d320e43b2ea5b97d08 Mon Sep 17 00:00:00 2001 From: Xianjie <5410381+qiaoxj07@users.noreply.github.com> Date: Mon, 31 Aug 2026 23:11:02 +0800 Subject: [PATCH 1/2] [TRTLLM-13409][fix] Count async KV completions as benchmark fill progress Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 58 ++++++++++++------- .../_torch/executor/test_benchmark_disagg.py | 55 ++++++++++++++++-- .../_torch/executor/test_py_executor.py | 6 +- 3 files changed, 89 insertions(+), 30 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 416e1a3721b3..8dd1972a5dc4 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -930,11 +930,11 @@ def on_detected(): # normal dummy add-forward-terminate lifecycle handles taper-down. # Only relevant in benchmark disagg mode; False otherwise. self._benchmark_fill_phase_active = self.is_benchmark_disagg - # Set when a blocking generation transfer returns during benchmark - # fill. The gate consumes this signal to retry immediately instead of - # adding an unnecessary polling delay after synchronous progress. - self._sync_disagg_transfer_made_progress = False - self._benchmark_sync_progress_global = False + # Set when a generation transfer completes during benchmark fill. The + # gate consumes this signal to retry immediately instead of adding an + # unnecessary polling delay after transfer progress. + self._disagg_gen_transfer_made_progress = False + self._benchmark_transfer_progress_global = False # Slow-start admission cap for benchmark disagg fill (see # _pop_from_waiting_queue). 0 = uninitialised; first throttled iter # seeds it to tp_size and each subsequent iter doubles it. @@ -3729,7 +3729,7 @@ def _sync_gen_only_benchmark_has_insufficient_kv( return all_ranks_fetched and any_rank_terminal_no_fit def _prepare_and_schedule_batch(self): - self._sync_disagg_transfer_made_progress = False + self._disagg_gen_transfer_made_progress = False self._poll_encoder_steps() new_requests = self._fetch_and_activate_new_requests() if self.should_stop_processing: @@ -3898,7 +3898,7 @@ def _kv_connector_wait_for_save(self): def _is_benchmark_disagg_fill_complete( self, scheduled_batch: ScheduledRequests, - local_sync_progress: bool = False) -> bool: + local_transfer_progress: bool = False) -> bool: """State-based fill-complete predicate for benchmark disagg mode. The gate opens when all three conditions hold globally: @@ -3909,7 +3909,7 @@ def _is_benchmark_disagg_fill_complete( KV-transfer phase (not in INIT, TRANS_IN_PROGRESS, or ERROR). (C) The KV cache transceiver has no pending receive sessions. - The conditions and synchronous-progress signal are gathered across the + The conditions and transfer-progress signal are gathered across the TP+CP scheduling group so every model-parallel rank makes the same gate and sleep decision. @@ -3918,8 +3918,8 @@ def _is_benchmark_disagg_fill_complete( Args: scheduled_batch: Passed for API compatibility with callers but no longer used by this predicate. - local_sync_progress: Whether this rank completed a synchronous KV - transfer in the current iteration. + local_transfer_progress: Whether this rank completed a KV transfer + in the current iteration. Returns: True when the fill phase is complete and the first forward @@ -3957,12 +3957,12 @@ def _is_benchmark_disagg_fill_complete( local_ok = int(local_all_fetched and local_all_past_transfer and local_no_inflight) - local_status = (local_ok, bool(local_sync_progress)) + local_status = (local_ok, bool(local_transfer_progress)) all_rank_status = self._allgather_model_parallel_status(local_status) all_ranks_ok = [status[0] for status in all_rank_status] global_ok = min(all_ranks_ok) == 1 - self._benchmark_sync_progress_global = any( + self._benchmark_transfer_progress_global = any( status[1] for status in all_rank_status) if self.dist.rank == 0: @@ -4001,8 +4001,8 @@ def _check_benchmark_disagg_gate(self, scheduled_batch: ScheduledRequests, A short sleep (0.1s) yields the CPU between retries that made no transfer progress while keeping the polling interval short enough to - avoid KV transfer backpressure on the CTX server. Synchronous receives - already block until they make progress, so those retries do not sleep. + avoid KV transfer backpressure on the CTX server. Retries that complete + a transfer do not sleep. Args: scheduled_batch: The current scheduled batch. @@ -4013,20 +4013,19 @@ def _check_benchmark_disagg_gate(self, scheduled_batch: ScheduledRequests, the caller should ``continue`` to the next loop iteration. """ if not self.is_warmup and not can_forward: - sync_transfer_made_progress = getattr( - self, "_sync_disagg_transfer_made_progress", False) - self._sync_disagg_transfer_made_progress = False + transfer_made_progress = self._disagg_gen_transfer_made_progress + self._disagg_gen_transfer_made_progress = False can_forward = self._is_benchmark_disagg_fill_complete( - scheduled_batch, sync_transfer_made_progress) - sync_transfer_made_progress = self._benchmark_sync_progress_global + scheduled_batch, transfer_made_progress) + transfer_made_progress = self._benchmark_transfer_progress_global 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: + elif not transfer_made_progress: time.sleep(0.1) if not can_forward: - self._fail_if_fill_gate_stalled(sync_transfer_made_progress) + self._fail_if_fill_gate_stalled(transfer_made_progress) return can_forward, True return can_forward, False @@ -7468,7 +7467,7 @@ def _recv_disagg_gen_cache(self, new_gen_reqs): for req in new_gen_reqs: self.kv_cache_transceiver.request_and_receive_sync(req) if req.state == LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE: - self._sync_disagg_transfer_made_progress = True + self._disagg_gen_transfer_made_progress = True self._check_cache_transfer_errors("generation requests") return @@ -7635,8 +7634,23 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): @nvtx_range("_check_disagg_gen_cache_transfer_status") def _check_disagg_gen_cache_transfer_status(self, atLeastNum: int = 0): + # Python transceivers report completed IDs, while the C++ runtime only + # mutates request state. Capture both forms of completion so the fill + # watchdog remains independent of the selected transceiver runtime. + tracked_requests = () + is_benchmark_fill = (self.is_benchmark_disagg + and self._benchmark_fill_phase_active) + if is_benchmark_fill: + tracked_requests = tuple( + req for req in self.active_requests + if req.is_disagg_generation_transmission_in_progress) + gen_status = self.kv_cache_transceiver.check_gen_transfer_status( atLeastNum) + if (is_benchmark_fill and (gen_status.completed_request_ids or any( + req.state == LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE + for req in tracked_requests))): + self._disagg_gen_transfer_made_progress = True if gen_status.cancelled_requests: user_canceled_set = set(self.canceled_req_ids) for req in gen_status.cancelled_requests: diff --git a/tests/unittest/_torch/executor/test_benchmark_disagg.py b/tests/unittest/_torch/executor/test_benchmark_disagg.py index e4b1aa14a7c4..2476bbc00e91 100644 --- a/tests/unittest/_torch/executor/test_benchmark_disagg.py +++ b/tests/unittest/_torch/executor/test_benchmark_disagg.py @@ -30,6 +30,7 @@ import pytest +from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import GenTransferStatus from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState from tensorrt_llm._torch.pyexecutor.scheduler import RequestScheduler, ScheduledRequests @@ -99,8 +100,8 @@ def __init__( benchmark_req_queues_size > 0 and kv_cache_transceiver is not None ) self._benchmark_fill_phase_active = self.is_benchmark_disagg - self._sync_disagg_transfer_made_progress = False - self._benchmark_sync_progress_global = False + self._disagg_gen_transfer_made_progress = False + self._benchmark_transfer_progress_global = False self._fill_admit_cap = 0 self.enable_attention_dp = enable_attention_dp self.max_num_active_requests = max_num_active_requests @@ -130,6 +131,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 + _check_disagg_gen_cache_transfer_status = PyExecutor._check_disagg_gen_cache_transfer_status _fail_if_fill_gate_stalled = PyExecutor._fail_if_fill_gate_stalled @@ -460,7 +462,7 @@ def test_gate_retries_with_short_sleep_when_incomplete(self, mock_time): mock_time.sleep.assert_called_once_with(0.1) @patch("tensorrt_llm._torch.pyexecutor.py_executor.time") - def test_gate_retries_without_sleep_after_sync_transfer_progress(self, mock_time): + def test_gate_retries_without_sleep_after_transfer_progress(self, mock_time): reqs = [_make_active_request(in_init=True)] ex = MockBenchmarkExecutor( benchmark_req_queues_size=4, @@ -468,15 +470,58 @@ def test_gate_retries_without_sleep_after_sync_transfer_progress(self, mock_time num_fetch_requests=2, active_requests=reqs, ) - ex._sync_disagg_transfer_made_progress = True + ex._disagg_gen_transfer_made_progress = True can_forward, should_retry = ex._check_benchmark_disagg_gate(ScheduledRequests(), False) assert can_forward is False assert should_retry is True - assert ex._sync_disagg_transfer_made_progress is False + assert ex._disagg_gen_transfer_made_progress is False mock_time.sleep.assert_not_called() + @pytest.mark.parametrize( + "completed_request_ids", + [ + pytest.param([17], id="python_runtime"), + pytest.param([], id="cpp_runtime"), + ], + ) + @patch("tensorrt_llm._torch.pyexecutor.py_executor.time") + def test_async_completion_resets_fill_stall_watchdog(self, mock_time, completed_request_ids): + completed_req = _make_active_request(in_transfer=True) + completed_req.py_request_id = 17 + blocked_req = _make_active_request(in_init=True) + transceiver = _make_transceiver(transfer_complete=False) + + def complete_request(_at_least_num): + completed_req.state = LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE + completed_req.is_disagg_generation_transmission_in_progress = False + return GenTransferStatus(completed_request_ids, [], []) + + transceiver.check_gen_transfer_status.side_effect = complete_request + ex = MockBenchmarkExecutor( + benchmark_req_queues_size=2, + kv_cache_transceiver=transceiver, + num_fetch_requests=2, + active_requests=[completed_req, blocked_req], + ) + ex._benchmark_fill_stall_timeout_sec = 5.0 + ex._benchmark_fill_stall_since = 1000.0 + ex.canceled_req_ids = [] + ex._is_disagg_inflight_cancel_active = Mock(return_value=False) + ex._check_cache_transfer_errors = Mock() + mock_time.monotonic.return_value = 1006.0 + + ex._check_disagg_gen_cache_transfer_status(0) + can_forward, should_retry = ex._check_benchmark_disagg_gate(ScheduledRequests(), False) + + assert can_forward is False + assert should_retry is True + assert ex._benchmark_fill_stall_since is None + assert ex._disagg_gen_transfer_made_progress is False + mock_time.sleep.assert_not_called() + transceiver.check_gen_transfer_status.assert_called_once_with(0) + @patch("tensorrt_llm._torch.pyexecutor.py_executor.time") def test_gate_skips_sleep_on_all_adp_ranks_when_peer_makes_progress(self, mock_time): reqs = [_make_active_request(in_init=True)] diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 607f2b0871bc..cad226ec4ce8 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -1081,7 +1081,7 @@ def test_sync_receive_does_not_poll_async_status(self, monkeypatch): executor.kv_cache_transceiver.request_and_receive_async.assert_not_called() executor._check_disagg_gen_cache_transfer_status.assert_not_called() executor._check_cache_transfer_errors.assert_called_once_with("generation requests") - assert executor._sync_disagg_transfer_made_progress + assert executor._disagg_gen_transfer_made_progress def test_sync_receive_drains_batch_before_rank_aligned_error_vote(self, monkeypatch): monkeypatch.setenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP", "1") @@ -1098,7 +1098,7 @@ def test_sync_receive_drains_batch_before_rank_aligned_error_vote(self, monkeypa ] executor._handle_errors = Mock() executor._check_cache_transfer_errors = Mock() - executor._sync_disagg_transfer_made_progress = False + executor._disagg_gen_transfer_made_progress = False error_request = Mock( py_request_id=1, state=LlmRequestState.DISAGG_GENERATION_INIT, @@ -1131,7 +1131,7 @@ def complete_or_error(req): executor.kv_cache_transceiver.cancel_request.assert_not_called() executor._handle_errors.assert_not_called() executor._check_cache_transfer_errors.assert_called_once_with("generation requests") - assert executor._sync_disagg_transfer_made_progress + assert executor._disagg_gen_transfer_made_progress PyExecutor._handle_disagg_cache_errors_synced(executor) From c8b8ccc4228a512e9ae4041ab3ed6b92e29b5df6 Mon Sep 17 00:00:00 2001 From: Xianjie <5410381+qiaoxj07@users.noreply.github.com> Date: Tue, 1 Sep 2026 19:31:40 +0800 Subject: [PATCH 2/2] [TRTLLM-13409][fix] Address async fill watchdog review Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 37 ++++-- .../_torch/executor/test_benchmark_disagg.py | 116 ++++++++++++++++-- 2 files changed, 138 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 8dd1972a5dc4..314f331cef58 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -51,6 +51,7 @@ from tensorrt_llm.tools.profiler.host_profile_tools.host_profiler import ( get_global_profiler, host_profiler_context) +from ..disaggregation.base.transfer import get_unique_rid from ..disaggregation.executor.admission import \ DisaggTransferAdmissionController from ..disaggregation.executor.pp_termination import DisaggPPTerminationHandler @@ -935,6 +936,10 @@ def on_detected(): # unnecessary polling delay after transfer progress. self._disagg_gen_transfer_made_progress = False self._benchmark_transfer_progress_global = False + # A completion can surface through a status ID or request-state + # mutation. Credit each logical request once across both signals so it + # cannot reset the benchmark fill watchdog more than once. + self._benchmark_completed_gen_transfer_ids: set[int] = set() # Slow-start admission cap for benchmark disagg fill (see # _pop_from_waiting_queue). 0 = uninitialised; first throttled iter # seeds it to tp_size and each subsequent iter doubles it. @@ -3598,7 +3603,7 @@ def _allgather_model_parallel_status( Args: local_status: Caller-defined ``(state, flag)`` pair from this rank. - The fill gate uses ``(ready, synchronous_progress)`` and the + The fill gate uses ``(ready, transfer_progress)`` and the fail-fast path uses ``(all_fetched, terminal_no_fit)``. Returns: @@ -4022,6 +4027,7 @@ def _check_benchmark_disagg_gate(self, scheduled_batch: ScheduledRequests, self._benchmark_fill_phase_active = False self._fill_admit_cap = 0 self._benchmark_fill_stall_since = None + self._benchmark_completed_gen_transfer_ids.clear() elif not transfer_made_progress: time.sleep(0.1) if not can_forward: @@ -7638,19 +7644,34 @@ def _check_disagg_gen_cache_transfer_status(self, atLeastNum: int = 0): # mutates request state. Capture both forms of completion so the fill # watchdog remains independent of the selected transceiver runtime. tracked_requests = () - is_benchmark_fill = (self.is_benchmark_disagg - and self._benchmark_fill_phase_active) - if is_benchmark_fill: + track_benchmark_fill_progress = (self.is_benchmark_disagg + and not self.is_warmup + and self._dist_size( + self.dist, "pp_size") == 1 + and self._benchmark_fill_phase_active) + if track_benchmark_fill_progress: tracked_requests = tuple( req for req in self.active_requests if req.is_disagg_generation_transmission_in_progress) gen_status = self.kv_cache_transceiver.check_gen_transfer_status( atLeastNum) - if (is_benchmark_fill and (gen_status.completed_request_ids or any( - req.state == LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE - for req in tracked_requests))): - self._disagg_gen_transfer_made_progress = True + if track_benchmark_fill_progress: + completed_request_ids = set(gen_status.completed_request_ids) + for req in tracked_requests: + if (req.state + != LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE): + continue + request_id = get_unique_rid(req) + if request_id is not None: + completed_request_ids.add(request_id) + new_completed_request_ids = ( + completed_request_ids - + self._benchmark_completed_gen_transfer_ids) + if new_completed_request_ids: + self._benchmark_completed_gen_transfer_ids.update( + new_completed_request_ids) + self._disagg_gen_transfer_made_progress = True if gen_status.cancelled_requests: user_canceled_set = set(self.canceled_req_ids) for req in gen_status.cancelled_requests: diff --git a/tests/unittest/_torch/executor/test_benchmark_disagg.py b/tests/unittest/_torch/executor/test_benchmark_disagg.py index 2476bbc00e91..ca89f11502bc 100644 --- a/tests/unittest/_torch/executor/test_benchmark_disagg.py +++ b/tests/unittest/_torch/executor/test_benchmark_disagg.py @@ -88,6 +88,7 @@ def __init__( kv_cache_transceiver=None, enable_attention_dp: bool = False, tp_size: int = 1, + pp_size: int = 1, rank: int = 0, num_fetch_requests: int = 0, is_warmup: bool = False, @@ -102,6 +103,7 @@ def __init__( self._benchmark_fill_phase_active = self.is_benchmark_disagg self._disagg_gen_transfer_made_progress = False self._benchmark_transfer_progress_global = False + self._benchmark_completed_gen_transfer_ids = set() self._fill_admit_cap = 0 self.enable_attention_dp = enable_attention_dp self.max_num_active_requests = max_num_active_requests @@ -112,6 +114,7 @@ def __init__( self.dist = Mock() self.dist.rank = rank self.dist.tp_size = tp_size + self.dist.pp_size = pp_size self.dist.world_size = tp_size # State the gate's stall bound reads. 0 disables the bound, which is @@ -439,11 +442,13 @@ def test_gate_opens_when_fill_complete(self, mock_time): active_requests=reqs, ) assert ex._benchmark_fill_phase_active is True + ex._benchmark_completed_gen_transfer_ids = {17} can_forward, should_retry = ex._check_benchmark_disagg_gate(ScheduledRequests(), False) assert can_forward is True assert should_retry is False assert ex._benchmark_fill_phase_active is False + assert ex._benchmark_completed_gen_transfer_ids == set() mock_time.sleep.assert_not_called() @patch("tensorrt_llm._torch.pyexecutor.py_executor.time") @@ -480,23 +485,29 @@ def test_gate_retries_without_sleep_after_transfer_progress(self, mock_time): mock_time.sleep.assert_not_called() @pytest.mark.parametrize( - "completed_request_ids", + "completion_signal", [ - pytest.param([17], id="python_runtime"), - pytest.param([], id="cpp_runtime"), + pytest.param("completed_ids", id="python_completed_ids"), + pytest.param("request_state", id="cpp_request_state"), ], ) @patch("tensorrt_llm._torch.pyexecutor.py_executor.time") - def test_async_completion_resets_fill_stall_watchdog(self, mock_time, completed_request_ids): + def test_async_completion_resets_fill_stall_watchdog(self, mock_time, completion_signal): completed_req = _make_active_request(in_transfer=True) completed_req.py_request_id = 17 + completed_req.request_id = 17 + completed_req.py_disaggregated_params = Mock(disagg_request_id=117) blocked_req = _make_active_request(in_init=True) transceiver = _make_transceiver(transfer_complete=False) def complete_request(_at_least_num): - completed_req.state = LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE - completed_req.is_disagg_generation_transmission_in_progress = False - return GenTransferStatus(completed_request_ids, [], []) + if completion_signal == "request_state": + completed_req.state = LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE + completed_req.is_disagg_generation_transmission_in_progress = False + return GenTransferStatus([], [], []) + # Isolate the Python transceiver's completed-ID signal: leave the + # request state unchanged so this branch cannot pass via mutation. + return GenTransferStatus([117], [], []) transceiver.check_gen_transfer_status.side_effect = complete_request ex = MockBenchmarkExecutor( @@ -513,15 +524,106 @@ def complete_request(_at_least_num): mock_time.monotonic.return_value = 1006.0 ex._check_disagg_gen_cache_transfer_status(0) + assert ex._disagg_gen_transfer_made_progress is True can_forward, should_retry = ex._check_benchmark_disagg_gate(ScheduledRequests(), False) assert can_forward is False assert should_retry is True assert ex._benchmark_fill_stall_since is None assert ex._disagg_gen_transfer_made_progress is False + assert ex._benchmark_completed_gen_transfer_ids == {117} mock_time.sleep.assert_not_called() transceiver.check_gen_transfer_status.assert_called_once_with(0) + @patch("tensorrt_llm._torch.pyexecutor.py_executor.time") + def test_zero_async_completions_do_not_reset_fill_stall_watchdog(self, mock_time): + in_progress_req = _make_active_request(in_transfer=True) + in_progress_req.py_request_id = 17 + transceiver = _make_transceiver(transfer_complete=False) + transceiver.check_gen_transfer_status.return_value = GenTransferStatus([], [], []) + ex = MockBenchmarkExecutor( + benchmark_req_queues_size=1, + kv_cache_transceiver=transceiver, + num_fetch_requests=1, + active_requests=[in_progress_req], + ) + ex._benchmark_fill_stall_timeout_sec = 5.0 + ex._benchmark_fill_stall_since = 1000.0 + ex.canceled_req_ids = [] + ex._is_disagg_inflight_cancel_active = Mock(return_value=False) + ex._check_cache_transfer_errors = Mock() + mock_time.monotonic.return_value = 1006.0 + + ex._check_disagg_gen_cache_transfer_status(0) + assert ex._disagg_gen_transfer_made_progress is False + with pytest.raises(RuntimeError, match="made no progress for 6s"): + ex._check_benchmark_disagg_gate(ScheduledRequests(), False) + + assert ex._benchmark_completed_gen_transfer_ids == set() + assert ex._benchmark_fill_stall_since is None + assert ex._disagg_gen_transfer_made_progress is False + mock_time.sleep.assert_called_once_with(0.1) + transceiver.check_gen_transfer_status.assert_called_once_with(0) + + @patch("tensorrt_llm._torch.pyexecutor.py_executor.time") + def test_same_async_completion_is_counted_once_across_signals(self, mock_time): + in_progress_req = _make_active_request(in_transfer=True) + in_progress_req.py_request_id = 17 + in_progress_req.request_id = 17 + in_progress_req.py_disaggregated_params = Mock(disagg_request_id=117) + transceiver = _make_transceiver(transfer_complete=False) + + def report_completion(_at_least_num): + if transceiver.check_gen_transfer_status.call_count == 1: + return GenTransferStatus([117], [], []) + in_progress_req.state = LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE + in_progress_req.is_disagg_generation_transmission_in_progress = False + return GenTransferStatus([], [], []) + + transceiver.check_gen_transfer_status.side_effect = report_completion + ex = MockBenchmarkExecutor( + benchmark_req_queues_size=2, + kv_cache_transceiver=transceiver, + num_fetch_requests=1, + active_requests=[in_progress_req], + ) + ex._benchmark_fill_stall_timeout_sec = 5.0 + ex.canceled_req_ids = [] + ex._is_disagg_inflight_cancel_active = Mock(return_value=False) + ex._check_cache_transfer_errors = Mock() + + ex._check_disagg_gen_cache_transfer_status(0) + assert ex._disagg_gen_transfer_made_progress is True + ex._check_benchmark_disagg_gate(ScheduledRequests(), False) + + ex._benchmark_fill_stall_since = 1000.0 + mock_time.monotonic.return_value = 1006.0 + ex._check_disagg_gen_cache_transfer_status(0) + assert ex._disagg_gen_transfer_made_progress is False + with pytest.raises(RuntimeError, match="made no progress for 6s"): + ex._check_benchmark_disagg_gate(ScheduledRequests(), False) + + assert ex._benchmark_completed_gen_transfer_ids == {117} + mock_time.sleep.assert_called_once_with(0.1) + + def test_pp_loop_does_not_scan_or_record_unused_fill_progress(self): + transceiver = _make_transceiver(transfer_complete=False) + transceiver.check_gen_transfer_status.return_value = GenTransferStatus([17], [], []) + ex = MockBenchmarkExecutor( + benchmark_req_queues_size=1, + kv_cache_transceiver=transceiver, + pp_size=2, + active_requests=Mock(), + ) + ex.canceled_req_ids = [] + ex._is_disagg_inflight_cancel_active = Mock(return_value=False) + ex._check_cache_transfer_errors = Mock() + + ex._check_disagg_gen_cache_transfer_status(0) + + assert ex._disagg_gen_transfer_made_progress is False + assert ex._benchmark_completed_gen_transfer_ids == set() + @patch("tensorrt_llm._torch.pyexecutor.py_executor.time") def test_gate_skips_sleep_on_all_adp_ranks_when_peer_makes_progress(self, mock_time): reqs = [_make_active_request(in_init=True)]