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

tensorrt_llm/bench/benchmark/__init__.py

@@ -55,6 +55,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

@@ -128,6 +128,13 @@
     help="Number of requests to cap benchmark run at. Minimum between value and"
     "length of dataset.",
 )
+@optgroup.option(
+    "--duration",
+    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).",
+)
 @optgroup.option(
     "--warmup",
     type=int,
@@ -382,7 +389,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

@@ -194,6 +194,13 @@
     "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=click.IntRange(min=1),
+    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,
@@ -518,7 +525,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
@@ -30,8 +44,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
@@ -44,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):
@@ -62,11 +89,18 @@ 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()
         sampling_params = copy.copy(sampling_params)
         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}")
@@ -112,6 +146,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
@@ -122,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(
@@ -197,9 +240,28 @@ 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."""
+        # 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()
+
+                # 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:
+                            self._inbox.get_nowait()
+                        except asyncio.QueueEmpty:
+                            break
+                    drain_in_flight = True
+                    break
+
                 try:
                     request, sampling_params, post_proc_params = self._inbox.get_nowait(
                     )
@@ -216,12 +278,33 @@ 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...")
+            # 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..."
+                )
+                # 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.
@@ -340,7 +423,25 @@ async def async_benchmark(
     iteration_log_addr: str = None,
     modality: Optional[str] = None,
     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()
@@ -351,7 +452,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)
@@ -376,7 +478,12 @@ 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"
+            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.")

tests/integration/test_lists/test-db/l0_cpu_arm.yml

@@ -15,3 +15,4 @@ l0_cpu_arm:
   tests:
   - unittest/executor/test_rpc.py
   - unittest/others/test_http_utils_fail_fast.py
+  - unittest/llmapi/test_bench_async.py

tests/integration/test_lists/test-db/l0_cpu_x86.yml

@@ -16,3 +16,4 @@ l0_cpu_x86:
   - unittest/executor/test_rpc.py
   - unittest/others/test_http_utils_fail_fast.py
   - unittest/executor/test_multi_frontend_routing.py
+  - unittest/llmapi/test_bench_async.py