From abb477eb881c35ec52b7adc0a9139f50c8aa791d Mon Sep 17 00:00:00 2001 From: Kuo Wei Date: Thu, 23 Apr 2026 14:18:00 +0000 Subject: [PATCH 01/18] [None][feat] Add duration-based execution to benchmark Signed-off-by: Kuo Wei --- tensorrt_llm/bench/benchmark/__init__.py | 2 + tensorrt_llm/bench/benchmark/low_latency.py | 9 +- tensorrt_llm/bench/benchmark/throughput.py | 9 +- .../bench/benchmark/utils/asynchronous.py | 47 ++++++++-- tests/unittest/llmapi/test_bench_async.py | 85 +++++++++++++++++++ 5 files changed, 145 insertions(+), 7 deletions(-) create mode 100644 tests/unittest/llmapi/test_bench_async.py diff --git a/tensorrt_llm/bench/benchmark/__init__.py b/tensorrt_llm/bench/benchmark/__init__.py index 83fd3e066614..67028daf30a4 100644 --- a/tensorrt_llm/bench/benchmark/__init__.py +++ b/tensorrt_llm/bench/benchmark/__init__.py @@ -30,6 +30,8 @@ class GeneralExecSettings(BaseModel): validation_alias=AliasChoices( "dataset_path", "dataset"), description="Path to dataset file") + duration: Optional[int] = Field( + default=None, description="Maximum run time in seconds") engine_dir: Optional[Path] = Field( default=None, description="Path to a serialized TRT-LLM engine") eos_id: int = Field( diff --git a/tensorrt_llm/bench/benchmark/low_latency.py b/tensorrt_llm/bench/benchmark/low_latency.py index 48477b823e80..3ea47a6cfc9f 100644 --- a/tensorrt_llm/bench/benchmark/low_latency.py +++ b/tensorrt_llm/bench/benchmark/low_latency.py @@ -127,6 +127,12 @@ help="Number of requests to cap benchmark run at. Minimum between value and" "length of dataset.", ) +@optgroup.option( + "--duration", + type=int, + default=None, + help="Maximum run time in seconds. Benchmark stops at whichever limit is hit first (num_requests or duration).", +) @optgroup.option( "--warmup", type=int, @@ -375,7 +381,8 @@ def latency_command( True, options.concurrency, iteration_writer.full_address, - modality=options.modality)) + modality=options.modality, + duration=options.duration)) logger.info("Benchmark done. Reporting results...") diff --git a/tensorrt_llm/bench/benchmark/throughput.py b/tensorrt_llm/bench/benchmark/throughput.py index 96a477195dd6..f1c33cc0c165 100755 --- a/tensorrt_llm/bench/benchmark/throughput.py +++ b/tensorrt_llm/bench/benchmark/throughput.py @@ -193,6 +193,12 @@ "Number of requests to cap benchmark run at. If not specified or set to 0, it will be the " "length of dataset.", ) +@optgroup.option( + "--duration", + type=int, + default=None, + help="Maximum run time in seconds. Benchmark stops at whichever limit is hit first (num_requests or duration).", +) @optgroup.option( "--warmup", type=int, @@ -501,7 +507,8 @@ def throughput_command( options.concurrency, iteration_writer.full_address, modality=options.modality, - tokenizer=multi_turn_tokenizer)) + tokenizer=multi_turn_tokenizer, + duration=options.duration)) logger.info("Benchmark done. Reporting results...") if options.modality is not None: diff --git a/tensorrt_llm/bench/benchmark/utils/asynchronous.py b/tensorrt_llm/bench/benchmark/utils/asynchronous.py index 9d8c0ca70688..6ecbe1afaa96 100644 --- a/tensorrt_llm/bench/benchmark/utils/asynchronous.py +++ b/tensorrt_llm/bench/benchmark/utils/asynchronous.py @@ -1,3 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-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. from __future__ import annotations import asyncio @@ -31,8 +45,11 @@ def __init__(self, streaming: bool, concurrency: int = -1, modality: Optional[str] = None, - tokenizer: Optional[PreTrainedTokenizer] = None) -> None: + tokenizer: Optional[PreTrainedTokenizer] = None, + duration: Optional[int] = None) -> None: self.llm = llm + self.duration = duration + self.start_time: Optional[float] = None self._inbox: asyncio.Queue[Tuple[InferenceRequest, SamplingParams]] = asyncio.Queue() self._outbox = outbox @@ -64,6 +81,8 @@ async def _process_single_request(self, request: InferenceRequest, sampling_params: SamplingParams, post_proc_params: PostprocParams): self.request_seen.set() + if self.start_time is None: + self.start_time = time.perf_counter() sampling_params = copy.copy(sampling_params) sampling_params.max_tokens = request.output_tokens @@ -113,6 +132,8 @@ async def _process_multi_turn_request(self, request: InferenceRequest, slot so that the conversation history stays consistent. """ self.request_seen.set() + if self.start_time is None: + self.start_time = time.perf_counter() sampling_params = copy.copy(sampling_params) sampling_params.max_tokens = request.output_tokens tokenizer = self.tokenizer @@ -200,6 +221,16 @@ async def worker(self) -> None: try: while not self._stop.is_set(): self._raise_for_failed_tasks() + + if self.duration is not None and self.start_time and time.perf_counter() - self.start_time >= self.duration: + logger.info("Duration reached. Stopping pulling requests.") + while not self._inbox.empty(): + try: + self._inbox.get_nowait() + except asyncio.QueueEmpty: + break + break + try: request, sampling_params, post_proc_params = self._inbox.get_nowait( ) @@ -216,9 +247,13 @@ async def worker(self) -> None: except asyncio.CancelledError: logger.info("Worker task cancelled.") finally: - logger.debug("Worker task cancelling remaining requests...") - for task in self._tasks: - task.cancel() + logger.debug("Worker task finishing...") + if self._stop.is_set(): + logger.debug("Worker task cancelling remaining requests due to stop signal...") + for task in self._tasks: + task.cancel() + else: + logger.debug("Worker task exiting. Waiting for in-flight tasks to complete...") logger.debug("Waiting for requests...") if self._tasks: await asyncio.wait(self._tasks) @@ -340,6 +375,7 @@ async def async_benchmark( iteration_log_addr: str = None, modality: Optional[str] = None, tokenizer: Optional[PreTrainedTokenizer] = None, + duration: Optional[int] = None, ) -> StatsKeeper: outbox = asyncio.Queue() statistics = StatsKeeper() @@ -351,7 +387,8 @@ async def async_benchmark( streaming, concurrency=concurrency, modality=modality, - tokenizer=tokenizer) + tokenizer=tokenizer, + duration=duration) enqueue_task: Optional[asyncio.Task] = None try: backend.run(iteration_addr=iteration_log_addr) diff --git a/tests/unittest/llmapi/test_bench_async.py b/tests/unittest/llmapi/test_bench_async.py new file mode 100644 index 000000000000..3eb1b56154ce --- /dev/null +++ b/tests/unittest/llmapi/test_bench_async.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-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. + +import asyncio +import time +from unittest.mock import MagicMock + +import pytest + +from tensorrt_llm import SamplingParams +from tensorrt_llm._tensorrt_engine import LLM +from tensorrt_llm.bench.benchmark.utils.asynchronous import LlmManager +from tensorrt_llm.bench.dataclasses.general import InferenceRequest +from tensorrt_llm.executor.postproc_worker import PostprocParams + + +@pytest.mark.asyncio +async def test_llm_manager_duration(): + # Mock LLM + mock_llm = MagicMock(spec=LLM) + mock_llm.args = MagicMock() + mock_llm.args.parallel_config = MagicMock() + mock_llm.args.parallel_config.world_size = 1 + + # Mock generate_async to return a mock output + mock_output = MagicMock() + mock_output.prompt_token_ids = [1, 2, 3] + mock_output.outputs = [MagicMock(token_ids=[4, 5])] + mock_output.finished = True + mock_output.id = 1 + mock_output.decoding_iter = 1 + + # We need to mock aresult() which is an async method + async def mock_aresult(): + await asyncio.sleep(0.6) # Make it take time + return mock_output + + mock_output.aresult = mock_aresult + mock_llm.generate_async.return_value = mock_output + + outbox = asyncio.Queue() + + manager = LlmManager( + llm=mock_llm, + outbox=outbox, + streaming=False, + concurrency=1, + duration=1, # 1 second + ) + + req = InferenceRequest(input_ids=[1, 2, 3], output_tokens=10) + sampling_params = SamplingParams() + post_proc_params = PostprocParams() + + # Enqueue 3 requests. Each takes 0.6s. + # Total time if all processed: 1.8s. + # With duration=1, it should stop after processing 2 requests. + await manager.enqueue(req, sampling_params, post_proc_params) + await manager.enqueue(req, sampling_params, post_proc_params) + await manager.enqueue(req, sampling_params, post_proc_params) + + manager.run() + + # Wait for more than 1 second (e.g., 1.5s) + await asyncio.sleep(1.5) + + # The worker should have stopped and cleared the inbox. + assert manager._inbox.empty() + + # 2 requests should have been processed and put into outbox. + assert outbox.qsize() == 2 + + await manager.stop() From 00c7a534e57e5d3176bd8997c0f6aba2d83e854c Mon Sep 17 00:00:00 2001 From: Kuo Wei Date: Thu, 23 Apr 2026 14:35:51 +0000 Subject: [PATCH 02/18] [None][chore] Improve docstring coverage in asynchronous.py Signed-off-by: Kuo Wei --- .../bench/benchmark/utils/asynchronous.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tensorrt_llm/bench/benchmark/utils/asynchronous.py b/tensorrt_llm/bench/benchmark/utils/asynchronous.py index 6ecbe1afaa96..f4850e875670 100644 --- a/tensorrt_llm/bench/benchmark/utils/asynchronous.py +++ b/tensorrt_llm/bench/benchmark/utils/asynchronous.py @@ -80,6 +80,7 @@ async def process_request(self, request: InferenceRequest, async def _process_single_request(self, request: InferenceRequest, sampling_params: SamplingParams, post_proc_params: PostprocParams): + """Process a single inference request.""" self.request_seen.set() if self.start_time is None: self.start_time = time.perf_counter() @@ -218,6 +219,7 @@ def _task_done_callback(self, task: asyncio.Task): self._task_errors.append(error) async def worker(self) -> None: + """Worker task that pulls requests from inbox and processes them.""" try: while not self._stop.is_set(): self._raise_for_failed_tasks() @@ -377,6 +379,23 @@ async def async_benchmark( tokenizer: Optional[PreTrainedTokenizer] = None, duration: Optional[int] = None, ) -> StatsKeeper: + """Run an asynchronous benchmark. + + Args: + llm: The LLM instance to use. + sampling_params: Sampling parameters for generation. + post_proc_params: Post-processing parameters. + requests: List of inference requests. + streaming: Whether to use streaming mode. + concurrency: Maximum concurrency limit. + iteration_log_addr: Address for iteration logging. + modality: Modality of requests. + tokenizer: Tokenizer for multi-turn requests. + duration: Maximum run time in seconds. + + Returns: + StatsKeeper containing benchmark statistics. + """ outbox = asyncio.Queue() statistics = StatsKeeper() submit_finished = asyncio.Event() From 1d1ca51243ed61e31d559c4ea9730a0593d36c21 Mon Sep 17 00:00:00 2001 From: Kuo Wei Date: Mon, 4 May 2026 08:55:23 +0000 Subject: [PATCH 03/18] [None][chore] Remove unused import in test_bench_async.py Signed-off-by: Kuo Wei --- tests/unittest/llmapi/test_bench_async.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/unittest/llmapi/test_bench_async.py b/tests/unittest/llmapi/test_bench_async.py index 3eb1b56154ce..7f509cd415a4 100644 --- a/tests/unittest/llmapi/test_bench_async.py +++ b/tests/unittest/llmapi/test_bench_async.py @@ -14,7 +14,6 @@ # limitations under the License. import asyncio -import time from unittest.mock import MagicMock import pytest From 16c8c702b3c2990b56e3b0fd20e1c1e4a0a5a024 Mon Sep 17 00:00:00 2001 From: Kuo Wei Date: Mon, 4 May 2026 08:59:57 +0000 Subject: [PATCH 04/18] [None][fix] Allow partial request completion when duration is set in async_benchmark Signed-off-by: Kuo Wei --- .../bench/benchmark/utils/asynchronous.py | 3 +- tests/unittest/llmapi/test_bench_async.py | 105 ++++++++++++++++++ 2 files changed, 107 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/bench/benchmark/utils/asynchronous.py b/tensorrt_llm/bench/benchmark/utils/asynchronous.py index f4850e875670..f588f437fb0f 100644 --- a/tensorrt_llm/bench/benchmark/utils/asynchronous.py +++ b/tensorrt_llm/bench/benchmark/utils/asynchronous.py @@ -432,7 +432,8 @@ async def async_benchmark( except asyncio.TimeoutError: logger.debug("No items in queue. Continuing.") - assert finished_requests == len(requests), "Benchmark failed" + if duration is None: + assert finished_requests == len(requests), "Benchmark failed" statistics.set_energy(monitor.total_energy) logger.info("Benchmark complete.") diff --git a/tests/unittest/llmapi/test_bench_async.py b/tests/unittest/llmapi/test_bench_async.py index 7f509cd415a4..2cac185b9567 100644 --- a/tests/unittest/llmapi/test_bench_async.py +++ b/tests/unittest/llmapi/test_bench_async.py @@ -82,3 +82,108 @@ async def mock_aresult(): assert outbox.qsize() == 2 await manager.stop() + + +@pytest.mark.asyncio +async def test_llm_manager_duration_not_exceeded(): + # Mock LLM + mock_llm = MagicMock(spec=LLM) + mock_llm.args = MagicMock() + mock_llm.args.parallel_config = MagicMock() + mock_llm.args.parallel_config.world_size = 1 + + # Mock generate_async to return a mock output + mock_output = MagicMock() + mock_output.prompt_token_ids = [1, 2, 3] + mock_output.outputs = [MagicMock(token_ids=[4, 5])] + mock_output.finished = True + mock_output.id = 1 + mock_output.decoding_iter = 1 + + async def mock_aresult(): + await asyncio.sleep(0.6) + return mock_output + + mock_output.aresult = mock_aresult + mock_llm.generate_async.return_value = mock_output + + outbox = asyncio.Queue() + + manager = LlmManager( + llm=mock_llm, + outbox=outbox, + streaming=False, + concurrency=1, + duration=5, # 5 seconds, plenty of time + ) + + req = InferenceRequest(input_ids=[1, 2, 3], output_tokens=10) + sampling_params = SamplingParams() + post_proc_params = PostprocParams() + + # Enqueue 2 requests. Each takes 0.6s. + # Total time: 1.2s. + # With duration=5, all requests should be processed. + await manager.enqueue(req, sampling_params, post_proc_params) + await manager.enqueue(req, sampling_params, post_proc_params) + + manager.run() + + # Wait for them to complete + await asyncio.sleep(1.5) + + # All 2 requests should have been processed and put into outbox. + assert outbox.qsize() == 2 + assert manager._inbox.empty() + + await manager.stop() + + +@pytest.mark.asyncio +async def test_async_benchmark_duration(): + from unittest.mock import patch + from tensorrt_llm.bench.benchmark.utils.asynchronous import async_benchmark + + # Mock LLM + mock_llm = MagicMock(spec=LLM) + mock_llm.args = MagicMock() + mock_llm.args.parallel_config = MagicMock() + mock_llm.args.parallel_config.world_size = 1 + + # Mock generate_async to return a mock output + mock_output = MagicMock() + mock_output.prompt_token_ids = [1, 2, 3] + mock_output.outputs = [MagicMock(token_ids=[4, 5])] + mock_output.finished = True + mock_output.id = 1 + mock_output.decoding_iter = 1 + + async def mock_aresult(): + await asyncio.sleep(0.6) # Make it take time + return mock_output + + mock_output.aresult = mock_aresult + mock_llm.generate_async.return_value = mock_output + + req = InferenceRequest(input_ids=[1, 2, 3], output_tokens=10) + requests = [req, req, req] + + # Patch EnergyMonitor and tqdm so we don't depend on actual NVML / environment + with patch('tensorrt_llm.bench.benchmark.utils.asynchronous.EnergyMonitor') as mock_energy, \ + patch('tensorrt_llm.bench.benchmark.utils.asynchronous.tqdm.tqdm') as mock_tqdm: + + # Mock the context manager of EnergyMonitor + mock_energy.return_value.__enter__.return_value.total_energy = 100.0 + + stats = await async_benchmark( + llm=mock_llm, + sampling_params=SamplingParams(), + post_proc_params=PostprocParams(), + requests=requests, + streaming=False, + duration=1, # 1 second limit + ) + + # Out of 3 requests taking 0.6s each, only 2 should complete within 1s limit + # Since we fixed the assertion, this should now finish gracefully without AssertionError + assert len(stats.requests) == 2 From a18ac0648797e8da308e2a411de05b144579758d Mon Sep 17 00:00:00 2001 From: Kuo Wei Date: Mon, 4 May 2026 09:02:10 +0000 Subject: [PATCH 05/18] [None][chore] Add informative log when duration limit is reached Signed-off-by: Kuo Wei --- tensorrt_llm/bench/benchmark/utils/asynchronous.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tensorrt_llm/bench/benchmark/utils/asynchronous.py b/tensorrt_llm/bench/benchmark/utils/asynchronous.py index f588f437fb0f..8820081b5339 100644 --- a/tensorrt_llm/bench/benchmark/utils/asynchronous.py +++ b/tensorrt_llm/bench/benchmark/utils/asynchronous.py @@ -434,6 +434,8 @@ async def async_benchmark( if duration is None: assert finished_requests == len(requests), "Benchmark failed" + elif finished_requests < len(requests): + logger.info(f"Duration limit reached. Processed {finished_requests}/{len(requests)} requests.") statistics.set_energy(monitor.total_energy) logger.info("Benchmark complete.") From 46d92f7f66236bc1decc3085e7280a8a6e8211e8 Mon Sep 17 00:00:00 2001 From: Kuo Wei Date: Thu, 16 Jul 2026 08:16:35 +0000 Subject: [PATCH 06/18] [None][chore] Fix pre-commit formatting and lint issues Signed-off-by: Kuo Wei --- tensorrt_llm/bench/benchmark/__init__.py | 4 ++-- tensorrt_llm/bench/benchmark/low_latency.py | 3 ++- tensorrt_llm/bench/benchmark/throughput.py | 3 ++- .../bench/benchmark/utils/asynchronous.py | 15 +++++++++++---- tests/unittest/llmapi/test_bench_async.py | 8 +++++--- 5 files changed, 22 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/bench/benchmark/__init__.py b/tensorrt_llm/bench/benchmark/__init__.py index 082a363f04ea..4a3b0e83f4e5 100644 --- a/tensorrt_llm/bench/benchmark/__init__.py +++ b/tensorrt_llm/bench/benchmark/__init__.py @@ -56,8 +56,8 @@ class GeneralExecSettings(BaseModel): validation_alias=AliasChoices( "dataset_path", "dataset"), description="Path to dataset file") - duration: Optional[int] = Field( - default=None, description="Maximum run time in seconds") + duration: Optional[int] = Field(default=None, + description="Maximum run time in seconds") engine_dir: Optional[Path] = Field( default=None, description="Path to a serialized TRT-LLM engine") eos_id: int = Field( diff --git a/tensorrt_llm/bench/benchmark/low_latency.py b/tensorrt_llm/bench/benchmark/low_latency.py index 6f98f61e278a..3a64de1169a7 100644 --- a/tensorrt_llm/bench/benchmark/low_latency.py +++ b/tensorrt_llm/bench/benchmark/low_latency.py @@ -132,7 +132,8 @@ "--duration", type=int, default=None, - help="Maximum run time in seconds. Benchmark stops at whichever limit is hit first (num_requests or duration).", + help= + "Maximum run time in seconds. Benchmark stops at whichever limit is hit first (num_requests or duration).", ) @optgroup.option( "--warmup", diff --git a/tensorrt_llm/bench/benchmark/throughput.py b/tensorrt_llm/bench/benchmark/throughput.py index 2b3b636cf438..78df35995ae2 100755 --- a/tensorrt_llm/bench/benchmark/throughput.py +++ b/tensorrt_llm/bench/benchmark/throughput.py @@ -198,7 +198,8 @@ "--duration", type=int, default=None, - help="Maximum run time in seconds. Benchmark stops at whichever limit is hit first (num_requests or duration).", + help= + "Maximum run time in seconds. Benchmark stops at whichever limit is hit first (num_requests or duration).", ) @optgroup.option( "--warmup", diff --git a/tensorrt_llm/bench/benchmark/utils/asynchronous.py b/tensorrt_llm/bench/benchmark/utils/asynchronous.py index c72ed0489a19..ccae0c5708ff 100644 --- a/tensorrt_llm/bench/benchmark/utils/asynchronous.py +++ b/tensorrt_llm/bench/benchmark/utils/asynchronous.py @@ -225,7 +225,8 @@ async def worker(self) -> None: while not self._stop.is_set(): self._raise_for_failed_tasks() - if self.duration is not None and self.start_time and time.perf_counter() - self.start_time >= self.duration: + if self.duration is not None and self.start_time and time.perf_counter( + ) - self.start_time >= self.duration: logger.info("Duration reached. Stopping pulling requests.") while not self._inbox.empty(): try: @@ -252,11 +253,15 @@ async def worker(self) -> None: finally: logger.debug("Worker task finishing...") if self._stop.is_set(): - logger.debug("Worker task cancelling remaining requests due to stop signal...") + logger.debug( + "Worker task cancelling remaining requests due to stop signal..." + ) for task in self._tasks: task.cancel() else: - logger.debug("Worker task exiting. Waiting for in-flight tasks to complete...") + logger.debug( + "Worker task exiting. Waiting for in-flight tasks to complete..." + ) logger.debug("Waiting for requests...") if self._tasks: await asyncio.wait(self._tasks) @@ -436,7 +441,9 @@ async def async_benchmark( if duration is None: assert finished_requests == len(requests), "Benchmark failed" elif finished_requests < len(requests): - logger.info(f"Duration limit reached. Processed {finished_requests}/{len(requests)} requests.") + logger.info( + f"Duration limit reached. Processed {finished_requests}/{len(requests)} requests." + ) statistics.set_energy(monitor.total_energy) logger.info("Benchmark complete.") diff --git a/tests/unittest/llmapi/test_bench_async.py b/tests/unittest/llmapi/test_bench_async.py index 2cac185b9567..bb224d8afed7 100644 --- a/tests/unittest/llmapi/test_bench_async.py +++ b/tests/unittest/llmapi/test_bench_async.py @@ -142,6 +142,7 @@ async def mock_aresult(): @pytest.mark.asyncio async def test_async_benchmark_duration(): from unittest.mock import patch + from tensorrt_llm.bench.benchmark.utils.asynchronous import async_benchmark # Mock LLM @@ -169,9 +170,10 @@ async def mock_aresult(): requests = [req, req, req] # Patch EnergyMonitor and tqdm so we don't depend on actual NVML / environment - with patch('tensorrt_llm.bench.benchmark.utils.asynchronous.EnergyMonitor') as mock_energy, \ - patch('tensorrt_llm.bench.benchmark.utils.asynchronous.tqdm.tqdm') as mock_tqdm: - + with ( + patch("tensorrt_llm.bench.benchmark.utils.asynchronous.EnergyMonitor") as mock_energy, + patch("tensorrt_llm.bench.benchmark.utils.asynchronous.tqdm.tqdm"), + ): # Mock the context manager of EnergyMonitor mock_energy.return_value.__enter__.return_value.total_energy = 100.0 From fc1503249c4bd9dfd4bc029e42864a660a5f2ddf Mon Sep 17 00:00:00 2001 From: Kuo Wei Date: Wed, 22 Jul 2026 22:29:56 +0800 Subject: [PATCH 07/18] [None][fix] Enforce benchmark duration at execution time The worker dispatches the entire inbox in one synchronous burst (no await point on the dispatch path), so start_time is still None when the duration check runs and every request is dispatched and completed regardless of the configured duration. Enforce the deadline in the request coroutines instead, after a concurrency slot is acquired: requests past the deadline are skipped while in-flight requests drain to completion, keeping recorded statistics unbiased. The multi-turn path stops issuing further turns past the deadline. Warn when --duration is combined with unbounded concurrency, where no client-side checkpoint exists to enforce it. Rework the unit test to await full worker drain (it previously passed only due to a timing artifact) and add a regression test for the eager-dispatch scenario raised in review. Signed-off-by: Kuo Wei --- .../bench/benchmark/utils/asynchronous.py | 28 ++++++- tests/unittest/llmapi/test_bench_async.py | 78 ++++++++++++++++++- 2 files changed, 101 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/bench/benchmark/utils/asynchronous.py b/tensorrt_llm/bench/benchmark/utils/asynchronous.py index 56d9721533fb..2d27402c6701 100644 --- a/tensorrt_llm/bench/benchmark/utils/asynchronous.py +++ b/tensorrt_llm/bench/benchmark/utils/asynchronous.py @@ -61,11 +61,21 @@ def __init__(self, self._iteration_log_task: Optional[asyncio.Task] = None self._concurrency_semaphore = asyncio.Semaphore( concurrency) if concurrency > 0 else None + if duration is not None and self._concurrency_semaphore is None: + logger.warning( + "--duration requires a concurrency limit to take effect; " + "without one, all requests are submitted to the engine " + "immediately and the full dataset will run.") self.streaming = streaming self.request_seen = asyncio.Event() self.modality = modality self.tokenizer = tokenizer + def _duration_exceeded(self) -> bool: + """Return whether the duration limit has elapsed since the first request.""" + return (self.duration is not None and self.start_time is not None + and time.perf_counter() - self.start_time >= self.duration) + async def process_request(self, request: InferenceRequest, sampling_params: SamplingParams, post_proc_params: PostprocParams): @@ -87,6 +97,10 @@ async def _process_single_request(self, request: InferenceRequest, sampling_params.max_tokens = request.output_tokens async with semaphore_guard(self._concurrency_semaphore): + # The worker dispatches the whole inbox eagerly, so the duration + # limit must be enforced here, after a concurrency slot is acquired. + if self._duration_exceeded(): + return request_start_timestamp = time.perf_counter_ns() time_on_first_token = None logger.debug(f"request.lora_request: {request.lora_request}") @@ -144,11 +158,18 @@ async def _process_multi_turn_request(self, request: InferenceRequest, all_output_tokens: List[int] = [] async with semaphore_guard(self._concurrency_semaphore): + # Enforce the duration limit at execution time (see + # _process_single_request). + if self._duration_exceeded(): + return request_start_timestamp = time.perf_counter_ns() time_on_first_token = None last_response = None for turn_id, question in enumerate(request.turns): + # Completed turns are still recorded below. + if turn_id > 0 and self._duration_exceeded(): + break messages.append({"role": "user", "content": question}) input_ids = await loop.run_in_executor( @@ -224,8 +245,11 @@ async def worker(self) -> None: while not self._stop.is_set(): self._raise_for_failed_tasks() - if self.duration is not None and self.start_time and time.perf_counter( - ) - self.start_time >= self.duration: + # Dispatch below never awaits, so this cannot fire until the + # inbox is fully dispatched; enforcement happens in the request + # coroutines (see _process_single_request). This exits the idle + # loop, and the drain keeps `busy` able to reach False. + if self._duration_exceeded(): logger.info("Duration reached. Stopping pulling requests.") while not self._inbox.empty(): try: diff --git a/tests/unittest/llmapi/test_bench_async.py b/tests/unittest/llmapi/test_bench_async.py index bb224d8afed7..9edc4cd45031 100644 --- a/tests/unittest/llmapi/test_bench_async.py +++ b/tests/unittest/llmapi/test_bench_async.py @@ -72,18 +72,90 @@ async def mock_aresult(): manager.run() - # Wait for more than 1 second (e.g., 1.5s) - await asyncio.sleep(1.5) + # Wait for the worker to fully drain: it exits its loop when the duration + # elapses and then awaits all in-flight tasks. Asserting after full drain + # (rather than sampling mid-flight) ensures requests dispatched past the + # deadline were actually skipped, not merely still running. + await asyncio.wait_for(manager._backend_task, timeout=10) # The worker should have stopped and cleared the inbox. assert manager._inbox.empty() - # 2 requests should have been processed and put into outbox. + # Requests 1 and 2 complete (0.6s + 0.6s); request 3 acquires its + # concurrency slot at t=1.2s, past the 1s deadline, and must be skipped. assert outbox.qsize() == 2 await manager.stop() +@pytest.mark.asyncio +async def test_llm_manager_duration_bounds_runtime_with_eager_dispatch(): + """Regression test for duration enforcement under eager task dispatch. + + The worker dispatches the entire inbox into tasks before any request + runs, so the duration must be enforced at execution time; otherwise + every dispatched request runs to completion and duration has no + effect on runtime. + """ + mock_llm = MagicMock(spec=LLM) + mock_llm.args = MagicMock() + mock_llm.args.parallel_config = MagicMock() + mock_llm.args.parallel_config.world_size = 1 + + mock_output = MagicMock() + mock_output.prompt_token_ids = [1, 2, 3] + mock_output.outputs = [MagicMock(token_ids=[4, 5])] + mock_output.finished = True + mock_output.id = 1 + mock_output.decoding_iter = 1 + + request_latency = 0.2 + + async def mock_aresult(): + await asyncio.sleep(request_latency) + return mock_output + + mock_output.aresult = mock_aresult + mock_llm.generate_async.return_value = mock_output + + outbox = asyncio.Queue() + num_requests = 20 + concurrency = 2 + duration = 1 + + manager = LlmManager( + llm=mock_llm, + outbox=outbox, + streaming=False, + concurrency=concurrency, + duration=duration, + ) + + req = InferenceRequest(input_ids=[1, 2, 3], output_tokens=10) + sampling_params = SamplingParams() + post_proc_params = PostprocParams() + + for _ in range(num_requests): + await manager.enqueue(req, sampling_params, post_proc_params) + + start = asyncio.get_running_loop().time() + manager.run() + await asyncio.wait_for(manager._backend_task, timeout=30) + elapsed = asyncio.get_running_loop().time() - start + + # Unbounded behavior would process all 20 requests in ~2s + # (20 / 2 slots * 0.2s). With enforcement, roughly + # duration / request_latency * concurrency = 10 requests complete. + # Generous bounds to stay robust on loaded CI machines. + assert outbox.qsize() < num_requests, ( + "Duration limit had no effect: all requests were processed." + ) + # Wall time is duration plus at most one in-flight drain, with slack. + assert elapsed < duration + request_latency + 0.5 + + await manager.stop() + + @pytest.mark.asyncio async def test_llm_manager_duration_not_exceeded(): # Mock LLM From acd262b6bc30b92a010215f1450338766b07b1d7 Mon Sep 17 00:00:00 2001 From: Kuo Wei Date: Thu, 23 Jul 2026 22:52:09 +0800 Subject: [PATCH 08/18] [None][test] Set concurrency=1 in test_async_benchmark_duration The test asserts that 2 of 3 requests (0.6s each) complete within a 1s duration, which requires requests to run back-to-back. Without a concurrency limit all three start immediately and finish before the deadline, so the assertion could not hold: duration cannot bound an unbounded-concurrency run, as there is no client-side checkpoint after dispatch. Pass concurrency=1, matching the test's sibling LlmManager duration tests and its own stated intent. Signed-off-by: Kuo Wei --- tests/unittest/llmapi/test_bench_async.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/unittest/llmapi/test_bench_async.py b/tests/unittest/llmapi/test_bench_async.py index 9edc4cd45031..9d978dc125b9 100644 --- a/tests/unittest/llmapi/test_bench_async.py +++ b/tests/unittest/llmapi/test_bench_async.py @@ -255,9 +255,13 @@ async def mock_aresult(): post_proc_params=PostprocParams(), requests=requests, streaming=False, + concurrency=1, duration=1, # 1 second limit ) - # Out of 3 requests taking 0.6s each, only 2 should complete within 1s limit - # Since we fixed the assertion, this should now finish gracefully without AssertionError + # With concurrency=1, requests run back-to-back (0.6s each): requests 1 + # and 2 complete at 0.6s and 1.2s; request 3 acquires its slot past the + # 1s deadline and is skipped. Without a concurrency limit all three would + # start immediately and finish within 0.6s, before the deadline — duration + # cannot bound an unbounded-concurrency run (see LlmManager warning). assert len(stats.requests) == 2 From 4a411b0ebcbe494ba85c04cbcc0ea7823db03252 Mon Sep 17 00:00:00 2001 From: Kuo Wei Date: Mon, 27 Jul 2026 23:06:06 +0800 Subject: [PATCH 09/18] [None][test] Fix broken imports and assertions in test_bench_async This file is not referenced by any test list, so CI has never collected it and three defects went unnoticed since it was added: - tensorrt_llm._tensorrt_engine was removed in #15918 and reached this branch through a main merge, so the module import failed at collection time and no test in the file could run. - InferenceRequest.task_id is required with no default, so every request construction raised ValidationError. - StatsKeeper keys its records by request id while the mock returned one shared id, so both completed requests collapsed into a single record. That merges their events and skews the recorded statistics, and it made the len(stats.requests) == 2 assertion unsatisfiable. Also harden two timing-sensitive waits. The duration-not-exceeded test slept a fixed 1.5s for work that takes 1.2s; it now awaits the perf items themselves, which is both faster and immune to scheduling delay. The eager-dispatch test now issues enough requests that an unenforced run would take ~5s against a ~1s enforced run, so the wall-clock bound has a wide margin in both directions instead of 0.3s. Signed-off-by: Kuo Wei --- tests/unittest/llmapi/test_bench_async.py | 80 +++++++++++++---------- 1 file changed, 47 insertions(+), 33 deletions(-) diff --git a/tests/unittest/llmapi/test_bench_async.py b/tests/unittest/llmapi/test_bench_async.py index 9d978dc125b9..6413226ff45e 100644 --- a/tests/unittest/llmapi/test_bench_async.py +++ b/tests/unittest/llmapi/test_bench_async.py @@ -14,12 +14,12 @@ # limitations under the License. import asyncio +import itertools from unittest.mock import MagicMock import pytest -from tensorrt_llm import SamplingParams -from tensorrt_llm._tensorrt_engine import LLM +from tensorrt_llm import LLM, SamplingParams from tensorrt_llm.bench.benchmark.utils.asynchronous import LlmManager from tensorrt_llm.bench.dataclasses.general import InferenceRequest from tensorrt_llm.executor.postproc_worker import PostprocParams @@ -59,7 +59,7 @@ async def mock_aresult(): duration=1, # 1 second ) - req = InferenceRequest(input_ids=[1, 2, 3], output_tokens=10) + req = InferenceRequest(task_id=0, input_ids=[1, 2, 3], output_tokens=10) sampling_params = SamplingParams() post_proc_params = PostprocParams() @@ -119,7 +119,10 @@ async def mock_aresult(): mock_llm.generate_async.return_value = mock_output outbox = asyncio.Queue() - num_requests = 20 + # Sized so that unenforced execution (num_requests / concurrency * + # request_latency = 5s) is far above the enforced runtime (~1s), leaving a + # wide margin for the wall-clock assertion below on a loaded machine. + num_requests = 50 concurrency = 2 duration = 1 @@ -131,7 +134,7 @@ async def mock_aresult(): duration=duration, ) - req = InferenceRequest(input_ids=[1, 2, 3], output_tokens=10) + req = InferenceRequest(task_id=0, input_ids=[1, 2, 3], output_tokens=10) sampling_params = SamplingParams() post_proc_params = PostprocParams() @@ -143,15 +146,15 @@ async def mock_aresult(): await asyncio.wait_for(manager._backend_task, timeout=30) elapsed = asyncio.get_running_loop().time() - start - # Unbounded behavior would process all 20 requests in ~2s - # (20 / 2 slots * 0.2s). With enforcement, roughly - # duration / request_latency * concurrency = 10 requests complete. - # Generous bounds to stay robust on loaded CI machines. + # With enforcement, roughly duration / request_latency * concurrency = 10 + # requests complete; without it, all 50 would. assert outbox.qsize() < num_requests, ( "Duration limit had no effect: all requests were processed." ) - # Wall time is duration plus at most one in-flight drain, with slack. - assert elapsed < duration + request_latency + 0.5 + # Wall time is duration plus at most one in-flight drain. The slack keeps + # this robust on a loaded machine while staying far below the ~5s an + # unenforced run would take. + assert elapsed < duration + request_latency + 1.5 await manager.stop() @@ -189,7 +192,7 @@ async def mock_aresult(): duration=5, # 5 seconds, plenty of time ) - req = InferenceRequest(input_ids=[1, 2, 3], output_tokens=10) + req = InferenceRequest(task_id=0, input_ids=[1, 2, 3], output_tokens=10) sampling_params = SamplingParams() post_proc_params = PostprocParams() @@ -201,11 +204,15 @@ async def mock_aresult(): manager.run() - # Wait for them to complete - await asyncio.sleep(1.5) + # Await the perf items themselves rather than sleeping a fixed interval. + # The worker loops until the 5s duration elapses, so there is no task + # completion to await here, and a fixed sleep would race the requests on a + # loaded machine. + for _ in range(2): + await asyncio.wait_for(outbox.get(), timeout=10) - # All 2 requests should have been processed and put into outbox. - assert outbox.qsize() == 2 + # Both requests were processed and none are left pending. + assert outbox.empty() assert manager._inbox.empty() await manager.stop() @@ -223,23 +230,30 @@ async def test_async_benchmark_duration(): mock_llm.args.parallel_config = MagicMock() mock_llm.args.parallel_config.world_size = 1 - # Mock generate_async to return a mock output - mock_output = MagicMock() - mock_output.prompt_token_ids = [1, 2, 3] - mock_output.outputs = [MagicMock(token_ids=[4, 5])] - mock_output.finished = True - mock_output.id = 1 - mock_output.decoding_iter = 1 - - async def mock_aresult(): - await asyncio.sleep(0.6) # Make it take time - return mock_output - - mock_output.aresult = mock_aresult - mock_llm.generate_async.return_value = mock_output - - req = InferenceRequest(input_ids=[1, 2, 3], output_tokens=10) - requests = [req, req, req] + # StatsKeeper records requests in a dict keyed by request id, so each + # response needs its own id -- a shared one would merge the requests into a + # single record and skew their timings. + def make_output(request_id): + output = MagicMock() + output.prompt_token_ids = [1, 2, 3] + output.outputs = [MagicMock(token_ids=[4, 5])] + output.finished = True + output.id = request_id + output.decoding_iter = 1 + + async def mock_aresult(): + await asyncio.sleep(0.6) # Make it take time + return output + + output.aresult = mock_aresult + return output + + response_ids = itertools.count() + mock_llm.generate_async.side_effect = lambda *args, **kwargs: make_output(next(response_ids)) + + requests = [ + InferenceRequest(task_id=i, input_ids=[1, 2, 3], output_tokens=10) for i in range(3) + ] # Patch EnergyMonitor and tqdm so we don't depend on actual NVML / environment with ( From 52ba2058bb44cae9f8e6c185fe180e613886b32d Mon Sep 17 00:00:00 2001 From: Kuo Wei Date: Mon, 27 Jul 2026 23:12:37 +0800 Subject: [PATCH 10/18] [None][test] Register test_bench_async in the CPU CI stages CI collects unit tests only through the lists under tests/integration/test_lists/test-db, so this file has never run since it was added and its failures went unnoticed. Register it in the two GPU-free pre-merge stages: the tests mock LLM entirely and need no device. Signed-off-by: Kuo Wei --- tests/integration/test_lists/test-db/l0_cpu_arm.yml | 1 + tests/integration/test_lists/test-db/l0_cpu_x86.yml | 1 + 2 files changed, 2 insertions(+) diff --git a/tests/integration/test_lists/test-db/l0_cpu_arm.yml b/tests/integration/test_lists/test-db/l0_cpu_arm.yml index ba693acaf756..16a3c8587767 100644 --- a/tests/integration/test_lists/test-db/l0_cpu_arm.yml +++ b/tests/integration/test_lists/test-db/l0_cpu_arm.yml @@ -14,3 +14,4 @@ l0_cpu_arm: orchestrator: mpi tests: - unittest/executor/test_rpc.py + - unittest/llmapi/test_bench_async.py diff --git a/tests/integration/test_lists/test-db/l0_cpu_x86.yml b/tests/integration/test_lists/test-db/l0_cpu_x86.yml index 9a39993347b8..f9e1dbe1abc1 100644 --- a/tests/integration/test_lists/test-db/l0_cpu_x86.yml +++ b/tests/integration/test_lists/test-db/l0_cpu_x86.yml @@ -14,3 +14,4 @@ l0_cpu_x86: orchestrator: mpi tests: - unittest/executor/test_rpc.py + - unittest/llmapi/test_bench_async.py From f8e0c31a704ee12fe78f53c304bca030a678869c Mon Sep 17 00:00:00 2001 From: Kuo Wei Date: Tue, 28 Jul 2026 20:21:13 +0800 Subject: [PATCH 11/18] [None][fix] Cancel in-flight requests when a benchmark request fails Draining in-flight requests instead of cancelling them is only correct when the duration limit is reached, where the point is to record their statistics. The check used self._stop, which is only set by an explicit shutdown, so a request failure took the drain path too: the worker waited for every outstanding request before surfacing the error, making a failing run slower than it was before duration support was added. Track the duration-triggered exit explicitly and cancel in every other case, restoring the original abort behaviour on failure while keeping the drain semantics on expiry. Add a regression test covering the failure path, which previously had no coverage. Signed-off-by: Kuo Wei --- .../bench/benchmark/utils/asynchronous.py | 14 ++-- tests/unittest/llmapi/test_bench_async.py | 64 +++++++++++++++++++ 2 files changed, 72 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/bench/benchmark/utils/asynchronous.py b/tensorrt_llm/bench/benchmark/utils/asynchronous.py index 2d27402c6701..ec2ffa80a921 100644 --- a/tensorrt_llm/bench/benchmark/utils/asynchronous.py +++ b/tensorrt_llm/bench/benchmark/utils/asynchronous.py @@ -241,6 +241,9 @@ def _task_done_callback(self, task: asyncio.Task): async def worker(self) -> None: """Worker task that pulls requests from inbox and processes them.""" + # Only a duration-triggered exit lets in-flight requests finish; a stop + # signal or a failure cancels them, as the pre-duration code always did. + drain_in_flight = False try: while not self._stop.is_set(): self._raise_for_failed_tasks() @@ -256,6 +259,7 @@ async def worker(self) -> None: self._inbox.get_nowait() except asyncio.QueueEmpty: break + drain_in_flight = True break try: @@ -275,16 +279,14 @@ async def worker(self) -> None: logger.info("Worker task cancelled.") finally: logger.debug("Worker task finishing...") - if self._stop.is_set(): + if drain_in_flight: logger.debug( - "Worker task cancelling remaining requests due to stop signal..." + "Duration reached. Waiting for in-flight requests to complete..." ) + else: + logger.debug("Worker task cancelling remaining requests...") for task in self._tasks: task.cancel() - else: - logger.debug( - "Worker task exiting. Waiting for in-flight tasks to complete..." - ) logger.debug("Waiting for requests...") if self._tasks: await asyncio.wait(self._tasks) diff --git a/tests/unittest/llmapi/test_bench_async.py b/tests/unittest/llmapi/test_bench_async.py index 6413226ff45e..dfd53cc04f62 100644 --- a/tests/unittest/llmapi/test_bench_async.py +++ b/tests/unittest/llmapi/test_bench_async.py @@ -159,6 +159,70 @@ async def mock_aresult(): await manager.stop() +@pytest.mark.asyncio +async def test_llm_manager_cancels_in_flight_requests_on_failure(): + """A failed request aborts the run instead of draining the in-flight ones. + + Only a duration-triggered exit waits for in-flight requests, so that their + statistics are recorded. On failure the benchmark is aborting anyway, and + waiting would make an erroring run slower than a successful one. + """ + mock_llm = MagicMock(spec=LLM) + mock_llm.args = MagicMock() + mock_llm.args.parallel_config = MagicMock() + mock_llm.args.parallel_config.world_size = 1 + + slow_latency = 30 # Far longer than the test should ever wait. + call_count = itertools.count() + + def generate_async(*args, **kwargs): + output = MagicMock() + output.prompt_token_ids = [1, 2, 3] + output.outputs = [MagicMock(token_ids=[4, 5])] + output.finished = True + output.id = next(call_count) + output.decoding_iter = 1 + + # The first request fails; the rest hang until cancelled. + if output.id == 0: + + async def mock_aresult(): + raise ValueError("simulated request failure") + else: + + async def mock_aresult(): + await asyncio.sleep(slow_latency) + return output + + output.aresult = mock_aresult + return output + + mock_llm.generate_async.side_effect = generate_async + + outbox = asyncio.Queue() + manager = LlmManager( + llm=mock_llm, + outbox=outbox, + streaming=False, + concurrency=4, + ) + + req = InferenceRequest(task_id=0, input_ids=[1, 2, 3], output_tokens=10) + for _ in range(4): + await manager.enqueue(req, SamplingParams(), PostprocParams()) + + start = asyncio.get_running_loop().time() + manager.run() + with pytest.raises(ValueError, match="simulated request failure"): + await asyncio.wait_for(manager._backend_task, timeout=10) + elapsed = asyncio.get_running_loop().time() - start + + # Without cancellation the worker would block on the 30s requests. + assert elapsed < slow_latency, ( + "Worker waited for in-flight requests instead of cancelling them." + ) + + @pytest.mark.asyncio async def test_llm_manager_duration_not_exceeded(): # Mock LLM From 77eb35480bd968a7d301c4556461341cf0e5f2eb Mon Sep 17 00:00:00 2001 From: Kuo Wei Date: Tue, 28 Jul 2026 20:21:13 +0800 Subject: [PATCH 12/18] [None][fix] Reject non-positive --duration values --duration accepted 0 and negative values. With 0 the deadline is already past once start_time is set for the first request, so every subsequent request is skipped and the benchmark produces almost no data while appearing to run normally. Constrain the option to positive integers so the mistake is reported at the command line instead. Leaving it unset still means no limit. Signed-off-by: Kuo Wei --- tensorrt_llm/bench/benchmark/low_latency.py | 2 +- tensorrt_llm/bench/benchmark/throughput.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/bench/benchmark/low_latency.py b/tensorrt_llm/bench/benchmark/low_latency.py index 3c9fb9752aa5..f19016e656aa 100644 --- a/tensorrt_llm/bench/benchmark/low_latency.py +++ b/tensorrt_llm/bench/benchmark/low_latency.py @@ -130,7 +130,7 @@ ) @optgroup.option( "--duration", - type=int, + type=click.IntRange(min=1), default=None, help= "Maximum run time in seconds. Benchmark stops at whichever limit is hit first (num_requests or duration).", diff --git a/tensorrt_llm/bench/benchmark/throughput.py b/tensorrt_llm/bench/benchmark/throughput.py index 1dcf18fde9e4..1e5a964f1ab5 100755 --- a/tensorrt_llm/bench/benchmark/throughput.py +++ b/tensorrt_llm/bench/benchmark/throughput.py @@ -196,7 +196,7 @@ ) @optgroup.option( "--duration", - type=int, + type=click.IntRange(min=1), default=None, help= "Maximum run time in seconds. Benchmark stops at whichever limit is hit first (num_requests or duration).", From af78879f9a1f331d6ed3a6f336429d904ea20301 Mon Sep 17 00:00:00 2001 From: Kuo Wei Date: Wed, 29 Jul 2026 23:03:50 +0800 Subject: [PATCH 13/18] [None][fix] Stop draining in-flight requests once one fails Draining on duration expiry waits for every in-flight request so their statistics are recorded. A failure during that wait was held back until the slowest sibling finished, so a hung request could delay the error indefinitely. Wait with FIRST_EXCEPTION instead and cancel the remainder as soon as a drained request raises. When every request succeeds the call still waits for all of them, leaving the duration-drain behaviour unchanged. Signed-off-by: Kuo Wei --- .../bench/benchmark/utils/asynchronous.py | 29 ++++++-- tests/unittest/llmapi/test_bench_async.py | 68 +++++++++++++++++++ 2 files changed, 90 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/bench/benchmark/utils/asynchronous.py b/tensorrt_llm/bench/benchmark/utils/asynchronous.py index ec2ffa80a921..9f7616b95e10 100644 --- a/tensorrt_llm/bench/benchmark/utils/asynchronous.py +++ b/tensorrt_llm/bench/benchmark/utils/asynchronous.py @@ -279,17 +279,32 @@ async def worker(self) -> None: logger.info("Worker task cancelled.") finally: logger.debug("Worker task finishing...") - if drain_in_flight: + # Snapshot: the done callback removes tasks from the set as they + # finish, and asyncio.wait must not race with that. + pending = set(self._tasks) + if not drain_in_flight: + logger.debug("Worker task cancelling remaining requests...") + for task in pending: + task.cancel() + elif pending: logger.debug( "Duration reached. Waiting for in-flight requests to complete..." ) - else: - logger.debug("Worker task cancelling remaining requests...") - for task in self._tasks: - task.cancel() + # Draining preserves the statistics of requests still running at + # the deadline, but a failure ends the run regardless, so stop + # draining as soon as one occurs. Without this a slow or hung + # sibling would hold the error back until it finished. + done, pending = await asyncio.wait( + pending, return_when=asyncio.FIRST_EXCEPTION) + if any(not task.cancelled() and task.exception() is not None + for task in done): + logger.debug( + "Request failed during drain. Cancelling the rest...") + for task in pending: + task.cancel() logger.debug("Waiting for requests...") - if self._tasks: - await asyncio.wait(self._tasks) + if pending: + await asyncio.wait(pending) self._raise_for_failed_tasks() # This asynchronous function acts as a worker that logs iteration statistics. diff --git a/tests/unittest/llmapi/test_bench_async.py b/tests/unittest/llmapi/test_bench_async.py index dfd53cc04f62..34fdfef32689 100644 --- a/tests/unittest/llmapi/test_bench_async.py +++ b/tests/unittest/llmapi/test_bench_async.py @@ -223,6 +223,74 @@ async def mock_aresult(): ) +@pytest.mark.asyncio +async def test_llm_manager_stops_draining_when_a_request_fails(): + """A failure during the duration drain is not held back by a hung request. + + Draining preserves the statistics of requests still running at the + deadline, but once one of them fails the run ends regardless, so the + remaining ones must not delay the error. + """ + mock_llm = MagicMock(spec=LLM) + mock_llm.args = MagicMock() + mock_llm.args.parallel_config = MagicMock() + mock_llm.args.parallel_config.world_size = 1 + + duration = 1 + fail_latency = 1.5 # Fails after the deadline, while draining. + hang_latency = 30 # Must never be waited on in full. + call_count = itertools.count() + + def generate_async(*args, **kwargs): + output = MagicMock() + output.prompt_token_ids = [1, 2, 3] + output.outputs = [MagicMock(token_ids=[4, 5])] + output.finished = True + output.id = next(call_count) + output.decoding_iter = 1 + + if output.id == 0: + + async def mock_aresult(): + await asyncio.sleep(fail_latency) + raise ValueError("simulated request failure") + else: + + async def mock_aresult(): + await asyncio.sleep(hang_latency) + return output + + output.aresult = mock_aresult + return output + + mock_llm.generate_async.side_effect = generate_async + + outbox = asyncio.Queue() + manager = LlmManager( + llm=mock_llm, + outbox=outbox, + streaming=False, + concurrency=2, + duration=duration, + ) + + req = InferenceRequest(task_id=0, input_ids=[1, 2, 3], output_tokens=10) + for _ in range(2): + await manager.enqueue(req, SamplingParams(), PostprocParams()) + + start = asyncio.get_running_loop().time() + manager.run() + with pytest.raises(ValueError, match="simulated request failure"): + await asyncio.wait_for(manager._backend_task, timeout=15) + elapsed = asyncio.get_running_loop().time() - start + + # The failure surfaces once it happens; draining for the hung request + # would delay it until hang_latency. + assert elapsed < hang_latency, ( + "Drain waited for the hung request instead of surfacing the failure." + ) + + @pytest.mark.asyncio async def test_llm_manager_duration_not_exceeded(): # Mock LLM From eea406011618f415472054a69bd983f61e1b2fd2 Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Wed, 29 Jul 2026 15:14:51 -0700 Subject: [PATCH 14/18] test: mark benchmark async tests as CPU-only Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- tests/unittest/llmapi/test_bench_async.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/unittest/llmapi/test_bench_async.py b/tests/unittest/llmapi/test_bench_async.py index 34fdfef32689..c037737dd577 100644 --- a/tests/unittest/llmapi/test_bench_async.py +++ b/tests/unittest/llmapi/test_bench_async.py @@ -24,6 +24,8 @@ from tensorrt_llm.bench.dataclasses.general import InferenceRequest from tensorrt_llm.executor.postproc_worker import PostprocParams +pytestmark = pytest.mark.cpu_only + @pytest.mark.asyncio async def test_llm_manager_duration(): From f3c5a111fcbe820e1be153a96894611ba681205a Mon Sep 17 00:00:00 2001 From: Kuo Wei Date: Thu, 30 Jul 2026 21:37:06 +0800 Subject: [PATCH 15/18] [None][fix] Require a concurrency limit when --duration is set throughput defaults concurrency to -1, meaning unlimited, and the deadline is applied after a request acquires its concurrency slot. With no limit there is no such slot, so every request reaches the engine at once and the run continues to the end of the dataset -- the flag silently did nothing on its most common invocation, contradicting its own help text. Reject the combination instead of warning about it. The check runs before the model is loaded so the mistake surfaces in seconds rather than after several minutes of startup. Signed-off-by: Kuo Wei --- tensorrt_llm/bench/benchmark/low_latency.py | 8 ++++++++ tensorrt_llm/bench/benchmark/throughput.py | 8 ++++++++ 2 files changed, 16 insertions(+) diff --git a/tensorrt_llm/bench/benchmark/low_latency.py b/tensorrt_llm/bench/benchmark/low_latency.py index f19016e656aa..24d87e253789 100644 --- a/tensorrt_llm/bench/benchmark/low_latency.py +++ b/tensorrt_llm/bench/benchmark/low_latency.py @@ -228,6 +228,14 @@ def latency_command( # Parameters from CLI # Model, experiment, and engine params options = get_general_cli_options(params, bench_env) + # Checked before the model is loaded so the mistake is reported in seconds + # rather than after several minutes of startup. + if options.duration is not None and options.concurrency <= 0: + raise click.UsageError( + "--duration requires a concurrency limit. Without one every request " + "is submitted to the engine at once, so there is no point at which " + "the deadline can be applied and the full dataset would run. Pass " + "--concurrency N.") # Speculative Decode Options medusa_choices = params.get("medusa_choices") diff --git a/tensorrt_llm/bench/benchmark/throughput.py b/tensorrt_llm/bench/benchmark/throughput.py index 1e5a964f1ab5..36faed83f3a1 100755 --- a/tensorrt_llm/bench/benchmark/throughput.py +++ b/tensorrt_llm/bench/benchmark/throughput.py @@ -335,6 +335,14 @@ def throughput_command( # Get general CLI options using the centralized function options: GeneralExecSettings = get_general_cli_options(params, bench_env) + # Checked before the model is loaded so the mistake is reported in seconds + # rather than after several minutes of startup. + if options.duration is not None and options.concurrency <= 0: + raise click.UsageError( + "--duration requires a concurrency limit. Without one every request " + "is submitted to the engine at once, so there is no point at which " + "the deadline can be applied and the full dataset would run. Pass " + "--concurrency N.") tokenizer = initialize_tokenizer(options.checkpoint_path, custom_tokenizer) # Extract throughput-specific options not handled by GeneralExecSettings From 28983fcef08cb0ffea837068bbdc7809e3c4497e Mon Sep 17 00:00:00 2001 From: Kuo Wei Date: Thu, 30 Jul 2026 21:38:17 +0800 Subject: [PATCH 16/18] [None][fix] Drop multi-turn requests truncated by the duration limit A multi-turn request that reaches the deadline stops issuing further turns but was still emitted as a completed perf item, so a conversation cut off after two of five turns was recorded alongside conversations that ran to completion. That mixes partial and full conversations in the same statistics. Drop the truncated request instead, matching how a single-turn request past the deadline is skipped without being recorded. Signed-off-by: Kuo Wei --- .../bench/benchmark/utils/asynchronous.py | 10 +++- tests/unittest/llmapi/test_bench_async.py | 60 +++++++++++++++++++ 2 files changed, 69 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/bench/benchmark/utils/asynchronous.py b/tensorrt_llm/bench/benchmark/utils/asynchronous.py index 9f7616b95e10..610727806b57 100644 --- a/tensorrt_llm/bench/benchmark/utils/asynchronous.py +++ b/tensorrt_llm/bench/benchmark/utils/asynchronous.py @@ -166,9 +166,10 @@ async def _process_multi_turn_request(self, request: InferenceRequest, time_on_first_token = None last_response = None + truncated = False for turn_id, question in enumerate(request.turns): - # Completed turns are still recorded below. if turn_id > 0 and self._duration_exceeded(): + truncated = True break messages.append({"role": "user", "content": question}) @@ -202,6 +203,13 @@ async def _process_multi_turn_request(self, request: InferenceRequest, last_response = response + if truncated: + # A conversation cut short by the deadline is not a completed + # request: recording it would mix partial and full conversations in + # the same statistics. Requests skipped before their first turn are + # dropped for the same reason. + return + response_end_timestamp = time.perf_counter_ns() request_perf_item = PerfItemTuple( diff --git a/tests/unittest/llmapi/test_bench_async.py b/tests/unittest/llmapi/test_bench_async.py index c037737dd577..258cfde1af5b 100644 --- a/tests/unittest/llmapi/test_bench_async.py +++ b/tests/unittest/llmapi/test_bench_async.py @@ -293,6 +293,66 @@ async def mock_aresult(): ) +@pytest.mark.asyncio +async def test_llm_manager_drops_truncated_multi_turn_request(): + """A conversation cut short by the deadline is not recorded. + + Its turns are incomplete, so emitting it would mix partial and full + conversations in the same statistics. + """ + mock_llm = MagicMock(spec=LLM) + mock_llm.args = MagicMock() + mock_llm.args.parallel_config = MagicMock() + mock_llm.args.parallel_config.world_size = 1 + + turn_latency = 0.6 + + def generate_async(*args, **kwargs): + output = MagicMock() + output.prompt_token_ids = [1, 2, 3] + output.outputs = [MagicMock(token_ids=[4, 5])] + output.finished = True + output.id = 1 + output.decoding_iter = 1 + + async def mock_aresult(): + await asyncio.sleep(turn_latency) + return output + + output.aresult = mock_aresult + return output + + mock_llm.generate_async.side_effect = generate_async + + tokenizer = MagicMock() + tokenizer.apply_chat_template.return_value = [1, 2, 3] + tokenizer.decode.return_value = "answer" + + outbox = asyncio.Queue() + manager = LlmManager( + llm=mock_llm, + outbox=outbox, + streaming=False, + concurrency=1, + duration=1, + tokenizer=tokenizer, + ) + + # Four turns at 0.6s each cannot finish inside the 1s deadline, so the + # conversation is cut short after turn 2. + req = InferenceRequest( + task_id=0, input_ids=[1, 2, 3], output_tokens=10, turns=["q1", "q2", "q3", "q4"] + ) + await manager.enqueue(req, SamplingParams(), PostprocParams()) + + manager.run() + await asyncio.wait_for(manager._backend_task, timeout=15) + + assert outbox.empty(), "A conversation truncated by the deadline was recorded as complete." + + await manager.stop() + + @pytest.mark.asyncio async def test_llm_manager_duration_not_exceeded(): # Mock LLM From ee8cf46c24aa66f51a1b0403c948b31338ab5df0 Mon Sep 17 00:00:00 2001 From: Kuo Wei Date: Thu, 30 Jul 2026 21:38:29 +0800 Subject: [PATCH 17/18] [None][test] Relax boundary assertions in the duration tests Both tests pin the exact number of requests that land on either side of the deadline: with a 1s limit and 0.6s requests, the second one acquires its slot at 0.6s and only just makes it. These now run in pre-merge on every PR, where a scheduling stall of a few hundred milliseconds pushes it past the deadline and fails a run that behaved correctly. Assert that the limit dropped requests rather than exactly how many crossed the boundary. A third request cannot complete either way, so the upper bound still catches a duration limit that has no effect. Also drop a duplicated worker log line that printed twice on a normal exit. Signed-off-by: Kuo Wei --- .../bench/benchmark/utils/asynchronous.py | 1 - tests/unittest/llmapi/test_bench_async.py | 22 +++++++++++-------- 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/bench/benchmark/utils/asynchronous.py b/tensorrt_llm/bench/benchmark/utils/asynchronous.py index 610727806b57..87b88e396d4a 100644 --- a/tensorrt_llm/bench/benchmark/utils/asynchronous.py +++ b/tensorrt_llm/bench/benchmark/utils/asynchronous.py @@ -282,7 +282,6 @@ async def worker(self) -> None: post_proc_params=post_proc_params)) task.add_done_callback(self._task_done_callback) self._tasks.add(task) - logger.debug("Worker task finishing...") except asyncio.CancelledError: logger.info("Worker task cancelled.") finally: diff --git a/tests/unittest/llmapi/test_bench_async.py b/tests/unittest/llmapi/test_bench_async.py index 258cfde1af5b..6a01047225dd 100644 --- a/tests/unittest/llmapi/test_bench_async.py +++ b/tests/unittest/llmapi/test_bench_async.py @@ -83,9 +83,12 @@ async def mock_aresult(): # The worker should have stopped and cleared the inbox. assert manager._inbox.empty() - # Requests 1 and 2 complete (0.6s + 0.6s); request 3 acquires its - # concurrency slot at t=1.2s, past the 1s deadline, and must be skipped. - assert outbox.qsize() == 2 + # Request 3 acquires its concurrency slot at t=1.2s, past the 1s deadline, + # so it is always skipped. Request 2 acquires at t=0.6s and normally makes + # the deadline, but a scheduling stall on a loaded machine can push it past; + # the point of the test is that the limit drops requests, not exactly how + # many land on the boundary. + assert 1 <= outbox.qsize() < 3 await manager.stop() @@ -467,9 +470,10 @@ async def mock_aresult(): duration=1, # 1 second limit ) - # With concurrency=1, requests run back-to-back (0.6s each): requests 1 - # and 2 complete at 0.6s and 1.2s; request 3 acquires its slot past the - # 1s deadline and is skipped. Without a concurrency limit all three would - # start immediately and finish within 0.6s, before the deadline — duration - # cannot bound an unbounded-concurrency run (see LlmManager warning). - assert len(stats.requests) == 2 + # With concurrency=1 requests run back-to-back (0.6s each), so request 3 + # acquires its slot past the 1s deadline and is always skipped. Request 2 + # sits on the boundary, so assert the limit took effect rather than the + # exact count. Without a concurrency limit all three would start at once + # and finish before the deadline, which is why the CLI rejects that + # combination. + assert 1 <= len(stats.requests) < 3 From 69ac657cd37c19355bf8d0340cc6618f4da10767 Mon Sep 17 00:00:00 2001 From: Kuo Wei Date: Fri, 31 Jul 2026 22:05:09 +0800 Subject: [PATCH 18/18] [None][chore] Document the reporting scope of --duration Requests dropped at the deadline never reach StatsKeeper, so a duration-bounded run reports throughput and latency over the requests that completed rather than the dataset that was submitted. Say so in the help text, along with the concurrency requirement. Signed-off-by: Kuo Wei --- tensorrt_llm/bench/benchmark/low_latency.py | 4 +++- tensorrt_llm/bench/benchmark/throughput.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/bench/benchmark/low_latency.py b/tensorrt_llm/bench/benchmark/low_latency.py index 24d87e253789..57755eee22b7 100644 --- a/tensorrt_llm/bench/benchmark/low_latency.py +++ b/tensorrt_llm/bench/benchmark/low_latency.py @@ -133,7 +133,9 @@ type=click.IntRange(min=1), default=None, help= - "Maximum run time in seconds. Benchmark stops at whichever limit is hit first (num_requests or duration).", + "Maximum run time in seconds. Benchmark stops at whichever limit is hit first (num_requests or duration). " + "Requests dropped at the deadline are excluded from the report, so the statistics cover the requests that " + "completed rather than the whole dataset.", ) @optgroup.option( "--warmup", diff --git a/tensorrt_llm/bench/benchmark/throughput.py b/tensorrt_llm/bench/benchmark/throughput.py index 36faed83f3a1..efbce5cedb26 100755 --- a/tensorrt_llm/bench/benchmark/throughput.py +++ b/tensorrt_llm/bench/benchmark/throughput.py @@ -199,7 +199,9 @@ type=click.IntRange(min=1), default=None, help= - "Maximum run time in seconds. Benchmark stops at whichever limit is hit first (num_requests or duration).", + "Maximum run time in seconds. Benchmark stops at whichever limit is hit first (num_requests or duration). " + "Requires --concurrency. Requests dropped at the deadline are excluded from the report, so the statistics " + "cover the requests that completed rather than the whole dataset.", ) @optgroup.option( "--warmup",