Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
Python
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155#!/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, andper-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`), plusstreaming (`stream: true`, `stream_options.include_usage: true`), which the Rust clientdoes not use but this bench needs for time-to-first-token. Reads `NEV_MODEL_URL` and`NEV_MODEL` from the environment (see the worktree's `.env`), each falling back to itsname from before the project was renamed (`GENERA_` in place of `NEV_`).
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 32Raw per-request records go to rig-throughput-results/<model>/concurrency-<N>.json (gitignored,per bench/.gitignore); one summary JSON line per level prints to stdout."""import jsonimport osimport statisticsimport sysimport timefrom concurrent.futures import ThreadPoolExecutorfrom pathlib import Path
import requests
BENCH_DIR = Path(__file__).resolve().parentRESULTS_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: """`NEV_<rest>`, else `GENERA_<rest>` -- the same fallback as `generation::env::var`.""" return os.environ.get(f"NEV_{rest}") or os.environ.get(f"GENERA_{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("NEV_MODEL_URL and NEV_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()