diff --git a/bench/exo_bench.py b/bench/exo_bench.py index 2daa4927..49922819 100644 --- a/bench/exo_bench.py +++ b/bench/exo_bench.py @@ -22,6 +22,7 @@ import contextlib import itertools import json import sys +import threading import time from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor, as_completed @@ -443,19 +444,37 @@ def main() -> int: all_rows.append(row) else: # Concurrent: fire N requests in parallel - # Each thread gets its own ExoClient (separate HTTP connection) + # Pre-build prompt once, barrier ensures simultaneous dispatch + content, actual_pp = prompt_sizer.build(pp) + pre_built_payload: dict[str, Any] = { + "model": full_model_id, + "messages": [{"role": "user", "content": content}], + "stream": False, + "max_tokens": tg, + } + barrier = threading.Barrier(concurrency) batch_results: list[tuple[dict[str, Any], int]] = [] batch_errors = 0 def _run_concurrent( - idx: int, *, _pp: int = pp, _tg: int = tg + idx: int, ) -> tuple[dict[str, Any], int]: c = ExoClient( args.host, args.port, timeout_s=args.timeout ) - return run_one_completion( - c, full_model_id, _pp, _tg, prompt_sizer - ) + barrier.wait() + t0 = time.perf_counter() + out = c.post_bench_chat_completions(pre_built_payload) + elapsed = time.perf_counter() - t0 + stats = out.get("generation_stats") + choices = out.get("choices") or [{}] + message = choices[0].get("message", {}) if choices else {} + text = message.get("content") or "" + return { + "elapsed_s": elapsed, + "output_text_preview": text[:200], + "stats": stats, + }, actual_pp with ThreadPoolExecutor(max_workers=concurrency) as pool: futures = {