#!/usr/bin/env python3 """Measures the rig's vLLM server against the real width prompt: output tokens/s, completions/s (the doc's "candidates per second"), time to first token, and per-request latency, at a set of concurrency levels. The request shape mirrors `crates/generation/src/model.rs`'s `HttpModelClient::complete` exactly (`chat_template_kwargs.enable_thinking: false`, same `max_tokens`), plus streaming (`stream: true`, `stream_options.include_usage: true`), which the Rust client does not use but this bench needs for time-to-first-token. Reads `WEBENGINE_MODEL_URL` and `WEBENGINE_MODEL` from the environment (see the worktree's `.env`). The prompt is read from `rig-throughput-results/prompt.txt`, the exact prompt `bin/real_model_run.rs` would assemble for its 8-candidate batch, written by `cargo run --release --example dump_prompt -p generation` (which needs `corpus/ground-truth` and `corpus/regression-set.json` already built -- see `crates/generation/src/bin/ground_truth_cache.rs` and `generation_loop.rs`). One round per concurrency level: N requests fired at once via a thread pool, streamed, all awaited before the next level starts. `completions_per_s` and `output_tokens_per_s` are fleet-wide rates over that round's wall time, not a mean of per-request rates. Usage: bench/rig_throughput.py [concurrency ...] # defaults to 1 4 8 16 32 Raw per-request records go to rig-throughput-results//concurrency-.json (gitignored, per bench/.gitignore); one summary JSON line per level prints to stdout. """ import json import os import statistics import sys import time from concurrent.futures import ThreadPoolExecutor from pathlib import Path import requests BENCH_DIR = Path(__file__).resolve().parent RESULTS_DIR = BENCH_DIR / "rig-throughput-results" PROMPT_PATH = RESULTS_DIR / "prompt.txt" # Matches `model.rs`'s `MAX_COMPLETION_TOKENS` -- the bench measures against the same # cap the real client sends, not an unbounded one. MAX_TOKENS = 3000 def project_env(rest: str) -> str | None: """`WEBENGINE_`, as `generation::env::var` reads it.""" return os.environ.get(f"WEBENGINE_{rest}") def one_request(base_url: str, model: str, prompt: str, seed: int) -> dict: """One streamed chat-completion request. `ttft_s` is wall-clock from request start to the first chunk carrying non-empty `delta.content` -- the model's first visible output token, not counting an empty role-only opening chunk some servers send.""" body = { "model": model, "messages": [{"role": "user", "content": prompt}], "chat_template_kwargs": {"enable_thinking": False}, "seed": seed, "max_tokens": MAX_TOKENS, "stream": True, "stream_options": {"include_usage": True}, } started = time.monotonic() ttft_s = None usage = None # Same bearer auth `HttpModelClient` sends when the key is set (the rig gateway needs it). key = project_env("MODEL_API_KEY") headers = {"Authorization": f"Bearer {key}"} if key else {} with requests.post( f"{base_url}/v1/chat/completions", json=body, headers=headers, stream=True, timeout=600 ) as resp: resp.raise_for_status() for line in resp.iter_lines(decode_unicode=True): if not line or not line.startswith("data: "): continue payload = line[len("data: "):] if payload == "[DONE]": break chunk = json.loads(payload) choices = chunk.get("choices") or [] if ttft_s is None and choices and choices[0].get("delta", {}).get("content"): ttft_s = time.monotonic() - started if chunk.get("usage"): usage = chunk["usage"] wall_s = time.monotonic() - started return { "seed": seed, "wall_s": wall_s, "ttft_s": ttft_s, "prompt_tokens": usage["prompt_tokens"] if usage else None, "completion_tokens": usage["completion_tokens"] if usage else None, } def model_results_dir(model: str) -> Path: """One raw-results directory per model, so measuring one model never overwrites another's per-request records.""" return RESULTS_DIR / model.replace("/", "_") def run_level(base_url: str, model: str, prompt: str, concurrency: int) -> dict: with ThreadPoolExecutor(max_workers=concurrency) as pool: started = time.monotonic() futures = [pool.submit(one_request, base_url, model, prompt, seed) for seed in range(concurrency)] records = [f.result() for f in futures] wall_s = time.monotonic() - started out_dir = model_results_dir(model) out_dir.mkdir(parents=True, exist_ok=True) with open(out_dir / f"concurrency-{concurrency}.json", "w") as f: json.dump(records, f, indent=2) total_completion_tokens = sum(r["completion_tokens"] or 0 for r in records) total_prompt_tokens = sum(r["prompt_tokens"] or 0 for r in records) latencies = [r["wall_s"] for r in records] ttfts = [r["ttft_s"] for r in records if r["ttft_s"] is not None] return { "concurrency": concurrency, "requests": len(records), "wall_s": wall_s, "completions_per_s": len(records) / wall_s, "output_tokens_per_s": total_completion_tokens / wall_s, "total_completion_tokens": total_completion_tokens, "total_prompt_tokens": total_prompt_tokens, "latency_s_mean": statistics.mean(latencies), "latency_s_median": statistics.median(latencies), "latency_s_max": max(latencies), "ttft_s_mean": statistics.mean(ttfts) if ttfts else None, "ttft_s_median": statistics.median(ttfts) if ttfts else None, } def main() -> None: base_url = project_env("MODEL_URL") model = project_env("MODEL") if not base_url or not model: sys.exit("WEBENGINE_MODEL_URL and WEBENGINE_MODEL must be set (see the worktree's .env)") if not PROMPT_PATH.exists(): sys.exit( f"{PROMPT_PATH} not found; run " "`cargo run --release --example dump_prompt -p generation` first " "(needs corpus/ground-truth and corpus/regression-set.json already built)" ) prompt = PROMPT_PATH.read_text() levels = [int(a) for a in sys.argv[1:]] or [1, 4, 8, 16, 32] for concurrency in levels: result = run_level(base_url, model, prompt, concurrency) print(json.dumps(result)) if __name__ == "__main__": main()