kanshan/TensorRT-LLMs
1
0
Fork 0
mirror of https://github.com/NVIDIA/TensorRT-LLM.git synced 2026-01-14 06:27:45 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

[feat] Add progress bar to benchmark (#5173)

Browse Source 
Signed-off-by: Rashid Kaleem <rkaleem@nvidia.com>
Co-authored-by: Zhihan Jiang <68881590+nvzhihanj@users.noreply.github.com>
Co-authored-by: Frank <3429989+FrankD412@users.noreply.github.com>
This commit is contained in:
Rashid Kaleem 2025-06-26 05:39:45 -05:00 committed by GitHub
parent 2eb6502b1d
commit 3a1f4d4001
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tensorrt_llm/bench/benchmark/utils/asynchronous.py
View File

@ -6,6 +6,7 @@ from contextlib import asynccontextmanager
from itertools import chain from itertools import chain
from typing import List, Optional, Set, Tuple from typing import List, Optional, Set, Tuple
import tqdm
from zmq import PUSH from zmq import PUSH
from zmq.asyncio import Context from zmq.asyncio import Context
@ -238,12 +239,14 @@ async def async_benchmark(
post_proc_params, submit_finished)) post_proc_params, submit_finished))
logger.info("Starting benchmark...") logger.info("Starting benchmark...")
pbar = tqdm.tqdm(total=len(requests), desc="Benchmarking")
while not submit_finished.is_set() or backend.busy or not outbox.empty( while not submit_finished.is_set() or backend.busy or not outbox.empty(
): ):
try: try:
item: PerfItemTuple = await asyncio.wait_for(outbox.get(), item: PerfItemTuple = await asyncio.wait_for(outbox.get(),
timeout=1.0) timeout=1.0)
statistics.register_request_perf_item(item) statistics.register_request_perf_item(item)
pbar.update(1)
except asyncio.TimeoutError: except asyncio.TimeoutError:
logger.debug("No items in queue. Continuing.") logger.debug("No items in queue. Continuing.")