mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
[feat] Add progress bar to benchmark (#5173)
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:
parent
2eb6502b1d
commit
3a1f4d4001
@ -6,6 +6,7 @@ from contextlib import asynccontextmanager
|
||||
from itertools import chain
|
||||
from typing import List, Optional, Set, Tuple
|
||||
|
||||
import tqdm
|
||||
from zmq import PUSH
|
||||
from zmq.asyncio import Context
|
||||
|
||||
@ -238,12 +239,14 @@ async def async_benchmark(
|
||||
post_proc_params, submit_finished))
|
||||
|
||||
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(
|
||||
):
|
||||
try:
|
||||
item: PerfItemTuple = await asyncio.wait_for(outbox.get(),
|
||||
timeout=1.0)
|
||||
statistics.register_request_perf_item(item)
|
||||
pbar.update(1)
|
||||
except asyncio.TimeoutError:
|
||||
logger.debug("No items in queue. Continuing.")
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user