[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:
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

View File

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