[None][feat] Add duration-based execution to benchmark

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(

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...")
 

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:

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)

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()