Skip to content

Commit ecdf13f

Browse files
committed
add tqdm
Signed-off-by: David Chen <[email protected]>
1 parent 10f0945 commit ecdf13f

File tree

4 files changed

+204
-17
lines changed

4 files changed

+204
-17
lines changed

benchmarks/online_serving/image_to_image/benchmark_i2i.py

Lines changed: 51 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
import time
1616
import urllib.error
1717
import urllib.request
18-
from concurrent.futures import ThreadPoolExecutor
18+
from concurrent.futures import ThreadPoolExecutor, as_completed
1919
from datetime import datetime, timezone
2020
from pathlib import Path
2121
from typing import Any, Optional
@@ -114,6 +114,27 @@ def _summary_stats(values: list[float]) -> dict[str, float]:
114114
}
115115

116116

117+
class _NullProgressBar:
118+
def update(self, n: int = 1) -> None:
119+
return
120+
121+
def close(self) -> None:
122+
return
123+
124+
125+
def _make_progress_bar(*, total: int, desc: str, enabled: bool) -> Any:
126+
if not enabled:
127+
return _NullProgressBar()
128+
try:
129+
from tqdm.auto import tqdm # type: ignore
130+
except Exception as e:
131+
raise RuntimeError(
132+
"tqdm is required for the progress bar. Install with `pip install tqdm` or pass --no-progress."
133+
) from e
134+
135+
return tqdm(total=total, desc=desc, unit="req", dynamic_ncols=True)
136+
137+
117138
def _join_url(base_url: str, path: str) -> str:
118139
base = base_url.rstrip("/")
119140
rel = path.lstrip("/")
@@ -330,6 +351,10 @@ def parse_args() -> argparse.Namespace:
330351
parser.add_argument("--output-json", type=str, default=None, help="Optional summary JSON output path.")
331352
parser.add_argument("--output-jsonl", type=str, default=None, help="Optional per-request JSONL output path.")
332353

354+
parser.add_argument("--progress", dest="progress", action="store_true", help="Enable tqdm progress bar.")
355+
parser.add_argument("--no-progress", dest="progress", action="store_false", help="Disable progress bar.")
356+
parser.set_defaults(progress=True)
357+
333358
return parser.parse_args()
334359

335360

@@ -416,9 +441,31 @@ def _task(i: int) -> dict[str, Any]:
416441
results: list[dict[str, Any]] = []
417442
start_benchmark = time.perf_counter()
418443
with ThreadPoolExecutor(max_workers=args.concurrency) as executor:
419-
futures = [executor.submit(_task, i) for i in range(args.num_requests)]
420-
for fut in futures:
421-
results.append(fut.result())
444+
future_to_idx = {executor.submit(_task, i): i for i in range(args.num_requests)}
445+
results_by_idx: list[Optional[dict[str, Any]]] = [None] * args.num_requests
446+
pbar = _make_progress_bar(
447+
total=args.num_requests,
448+
desc="Benchmark requests",
449+
enabled=args.progress,
450+
)
451+
try:
452+
for fut in as_completed(future_to_idx):
453+
idx = future_to_idx[fut]
454+
try:
455+
r = fut.result()
456+
except Exception as e:
457+
r = {
458+
"request_idx": idx,
459+
"latency_ms": float("nan"),
460+
"success": False,
461+
"error": str(e),
462+
"outputs": None,
463+
}
464+
results_by_idx[idx] = r
465+
pbar.update(1)
466+
finally:
467+
pbar.close()
468+
results = [r for r in results_by_idx if r is not None]
422469
total_time_s = time.perf_counter() - start_benchmark
423470

424471
if monitor is not None:
@@ -509,4 +556,3 @@ def _task(i: int) -> dict[str, Any]:
509556

510557
if __name__ == "__main__":
511558
main()
512-

benchmarks/online_serving/image_to_video/benchmark_i2v.py

Lines changed: 51 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
import time
1717
import urllib.error
1818
import urllib.request
19-
from concurrent.futures import ThreadPoolExecutor
19+
from concurrent.futures import ThreadPoolExecutor, as_completed
2020
from datetime import datetime, timezone
2121
from pathlib import Path
2222
from typing import Any, Optional
@@ -114,6 +114,27 @@ def _summary_stats(values: list[float]) -> dict[str, float]:
114114
}
115115

116116

117+
class _NullProgressBar:
118+
def update(self, n: int = 1) -> None:
119+
return
120+
121+
def close(self) -> None:
122+
return
123+
124+
125+
def _make_progress_bar(*, total: int, desc: str, enabled: bool) -> Any:
126+
if not enabled:
127+
return _NullProgressBar()
128+
try:
129+
from tqdm.auto import tqdm # type: ignore
130+
except Exception as e:
131+
raise RuntimeError(
132+
"tqdm is required for the progress bar. Install with `pip install tqdm` or pass --no-progress."
133+
) from e
134+
135+
return tqdm(total=total, desc=desc, unit="req", dynamic_ncols=True)
136+
137+
117138
def _join_url(base_url: str, path: str) -> str:
118139
base = base_url.rstrip("/")
119140
rel = path.lstrip("/")
@@ -382,6 +403,10 @@ def parse_args() -> argparse.Namespace:
382403
parser.add_argument("--output-json", type=str, default=None, help="Optional summary JSON output path.")
383404
parser.add_argument("--output-jsonl", type=str, default=None, help="Optional per-request JSONL output path.")
384405

406+
parser.add_argument("--progress", dest="progress", action="store_true", help="Enable tqdm progress bar.")
407+
parser.add_argument("--no-progress", dest="progress", action="store_false", help="Disable progress bar.")
408+
parser.set_defaults(progress=True)
409+
385410
return parser.parse_args()
386411

387412

@@ -469,9 +494,31 @@ def _task(i: int) -> dict[str, Any]:
469494
results: list[dict[str, Any]] = []
470495
start_benchmark = time.perf_counter()
471496
with ThreadPoolExecutor(max_workers=args.concurrency) as executor:
472-
futures = [executor.submit(_task, i) for i in range(args.num_requests)]
473-
for fut in futures:
474-
results.append(fut.result())
497+
future_to_idx = {executor.submit(_task, i): i for i in range(args.num_requests)}
498+
results_by_idx: list[Optional[dict[str, Any]]] = [None] * args.num_requests
499+
pbar = _make_progress_bar(
500+
total=args.num_requests,
501+
desc="Benchmark requests",
502+
enabled=args.progress,
503+
)
504+
try:
505+
for fut in as_completed(future_to_idx):
506+
idx = future_to_idx[fut]
507+
try:
508+
r = fut.result()
509+
except Exception as e:
510+
r = {
511+
"request_idx": idx,
512+
"latency_ms": float("nan"),
513+
"success": False,
514+
"error": str(e),
515+
"outputs": None,
516+
}
517+
results_by_idx[idx] = r
518+
pbar.update(1)
519+
finally:
520+
pbar.close()
521+
results = [r for r in results_by_idx if r is not None]
475522
total_time_s = time.perf_counter() - start_benchmark
476523

477524
if monitor is not None:

benchmarks/online_serving/text_to_image/benchmark_t2i.py

Lines changed: 51 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
import urllib.error
1515
import urllib.parse
1616
import urllib.request
17-
from concurrent.futures import ThreadPoolExecutor
17+
from concurrent.futures import ThreadPoolExecutor, as_completed
1818
from datetime import datetime, timezone
1919
from pathlib import Path
2020
from typing import Any, Optional
@@ -114,6 +114,27 @@ def _summary_stats(values: list[float]) -> dict[str, float]:
114114
}
115115

116116

117+
class _NullProgressBar:
118+
def update(self, n: int = 1) -> None:
119+
return
120+
121+
def close(self) -> None:
122+
return
123+
124+
125+
def _make_progress_bar(*, total: int, desc: str, enabled: bool) -> Any:
126+
if not enabled:
127+
return _NullProgressBar()
128+
try:
129+
from tqdm.auto import tqdm # type: ignore
130+
except Exception as e:
131+
raise RuntimeError(
132+
"tqdm is required for the progress bar. Install with `pip install tqdm` or pass --no-progress."
133+
) from e
134+
135+
return tqdm(total=total, desc=desc, unit="req", dynamic_ncols=True)
136+
137+
117138
def _join_url(base_url: str, path: str) -> str:
118139
base = base_url.rstrip("/")
119140
rel = path.lstrip("/")
@@ -318,6 +339,10 @@ def parse_args() -> argparse.Namespace:
318339
parser.add_argument("--output-json", type=str, default=None, help="Optional summary JSON output path.")
319340
parser.add_argument("--output-jsonl", type=str, default=None, help="Optional per-request JSONL output path.")
320341

342+
parser.add_argument("--progress", dest="progress", action="store_true", help="Enable tqdm progress bar.")
343+
parser.add_argument("--no-progress", dest="progress", action="store_false", help="Disable progress bar.")
344+
parser.set_defaults(progress=True)
345+
321346
return parser.parse_args()
322347

323348

@@ -401,9 +426,31 @@ def _task(i: int) -> dict[str, Any]:
401426
results: list[dict[str, Any]] = []
402427
start_benchmark = time.perf_counter()
403428
with ThreadPoolExecutor(max_workers=args.concurrency) as executor:
404-
futures = [executor.submit(_task, i) for i in range(args.num_requests)]
405-
for fut in futures:
406-
results.append(fut.result())
429+
future_to_idx = {executor.submit(_task, i): i for i in range(args.num_requests)}
430+
results_by_idx: list[Optional[dict[str, Any]]] = [None] * args.num_requests
431+
pbar = _make_progress_bar(
432+
total=args.num_requests,
433+
desc="Benchmark requests",
434+
enabled=args.progress,
435+
)
436+
try:
437+
for fut in as_completed(future_to_idx):
438+
idx = future_to_idx[fut]
439+
try:
440+
r = fut.result()
441+
except Exception as e:
442+
r = {
443+
"request_idx": idx,
444+
"latency_ms": float("nan"),
445+
"success": False,
446+
"error": str(e),
447+
"outputs": None,
448+
}
449+
results_by_idx[idx] = r
450+
pbar.update(1)
451+
finally:
452+
pbar.close()
453+
results = [r for r in results_by_idx if r is not None]
407454
total_time_s = time.perf_counter() - start_benchmark
408455

409456
if monitor is not None:

benchmarks/online_serving/text_to_video/benchmark_t2v.py

Lines changed: 51 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
import time
1414
import urllib.error
1515
import urllib.request
16-
from concurrent.futures import ThreadPoolExecutor
16+
from concurrent.futures import ThreadPoolExecutor, as_completed
1717
from datetime import datetime, timezone
1818
from pathlib import Path
1919
from typing import Any, Optional
@@ -112,6 +112,27 @@ def _summary_stats(values: list[float]) -> dict[str, float]:
112112
}
113113

114114

115+
class _NullProgressBar:
116+
def update(self, n: int = 1) -> None:
117+
return
118+
119+
def close(self) -> None:
120+
return
121+
122+
123+
def _make_progress_bar(*, total: int, desc: str, enabled: bool) -> Any:
124+
if not enabled:
125+
return _NullProgressBar()
126+
try:
127+
from tqdm.auto import tqdm # type: ignore
128+
except Exception as e:
129+
raise RuntimeError(
130+
"tqdm is required for the progress bar. Install with `pip install tqdm` or pass --no-progress."
131+
) from e
132+
133+
return tqdm(total=total, desc=desc, unit="req", dynamic_ncols=True)
134+
135+
115136
def _join_url(base_url: str, path: str) -> str:
116137
base = base_url.rstrip("/")
117138
rel = path.lstrip("/")
@@ -340,6 +361,10 @@ def parse_args() -> argparse.Namespace:
340361
parser.add_argument("--output-json", type=str, default=None, help="Optional summary JSON output path.")
341362
parser.add_argument("--output-jsonl", type=str, default=None, help="Optional per-request JSONL output path.")
342363

364+
parser.add_argument("--progress", dest="progress", action="store_true", help="Enable tqdm progress bar.")
365+
parser.add_argument("--no-progress", dest="progress", action="store_false", help="Disable progress bar.")
366+
parser.set_defaults(progress=True)
367+
343368
return parser.parse_args()
344369

345370

@@ -429,9 +454,31 @@ def _task(i: int) -> dict[str, Any]:
429454
results: list[dict[str, Any]] = []
430455
start_benchmark = time.perf_counter()
431456
with ThreadPoolExecutor(max_workers=args.concurrency) as executor:
432-
futures = [executor.submit(_task, i) for i in range(args.num_requests)]
433-
for fut in futures:
434-
results.append(fut.result())
457+
future_to_idx = {executor.submit(_task, i): i for i in range(args.num_requests)}
458+
results_by_idx: list[Optional[dict[str, Any]]] = [None] * args.num_requests
459+
pbar = _make_progress_bar(
460+
total=args.num_requests,
461+
desc="Benchmark requests",
462+
enabled=args.progress,
463+
)
464+
try:
465+
for fut in as_completed(future_to_idx):
466+
idx = future_to_idx[fut]
467+
try:
468+
r = fut.result()
469+
except Exception as e:
470+
r = {
471+
"request_idx": idx,
472+
"latency_ms": float("nan"),
473+
"success": False,
474+
"error": str(e),
475+
"outputs": None,
476+
}
477+
results_by_idx[idx] = r
478+
pbar.update(1)
479+
finally:
480+
pbar.close()
481+
results = [r for r in results_by_idx if r is not None]
435482
total_time_s = time.perf_counter() - start_benchmark
436483

437484
if monitor is not None:

0 commit comments

Comments
 (0)