#!/usr/bin/env python3 from __future__ import annotations import argparse import concurrent.futures import json import statistics import sys import time from dataclasses import asdict, dataclass from typing import Any from urllib.parse import urlparse import requests from datasets import get_dataset_config_names, load_dataset from tqdm import tqdm DATASET_REPO = "nvidia/SPEED-Bench" @dataclass class Sample: id: str category: str turns: list[str] @dataclass class RequestResult: id: str category: str ok: bool turns: int latency_s: float prompt_tokens: int completion_tokens: int total_tokens: int finish_reason: str | None draft_n: int draft_n_accepted: int prompt_ms: float | None predicted_ms: float | None prompt_per_second: float | None predicted_per_second: float | None error: str | None def normalize_base_url(url: str) -> str: url = url.strip().rstrip("/") if not url: raise ValueError("--url cannot be empty") if "://" not in url: url = "http://" + url parsed = urlparse(url) if not parsed.scheme or not parsed.netloc: raise ValueError(f"invalid --url: {url}") if not parsed.path.rstrip("/").endswith("/v1"): url = url + "/v1" return url.rstrip("/") def parse_extra_inputs(value: str) -> dict[str, Any]: extra = json.loads(value) if not isinstance(extra, dict): raise ValueError("--extra-inputs must be a JSON object") return extra def extract_turns(row: dict[str, Any]) -> list[str]: turns = row.get("turns") if isinstance(turns, list) and turns: clean_turns = [str(turn).strip() for turn in turns if turn and str(turn).strip()] if clean_turns: return clean_turns raise ValueError("missing or empty turns") def load_samples(args: argparse.Namespace) -> list[Sample]: bench_names = get_dataset_config_names(DATASET_REPO) if args.bench not in bench_names: raise ValueError( f"unknown --bench {args.bench!r}; available benches: {', '.join(bench_names)}" ) dataset = load_dataset(DATASET_REPO, name=args.bench, split="test") categories = list(dict.fromkeys(str(category) for category in dataset["category"])) requested_categories = None if args.category != "all": requested_list = [category.strip() for category in args.category.split(",") if category.strip()] if not requested_list: raise ValueError( f"--category must be 'all' or a comma-separated list; available categories: {', '.join(categories)}" ) requested_categories = set(requested_list) unknown_categories = [category for category in requested_list if category not in categories] if unknown_categories: unknown = ", ".join(unknown_categories) raise ValueError( f"unknown --category {unknown!r} for bench {args.bench!r}; " f"available categories: all, {', '.join(categories)}" ) samples: list[Sample] = [] samples_per_category: dict[str, int] = {} skipped = 0 for index, row_raw in enumerate(dataset): row = dict(row_raw) category_raw = row.get("category") if not isinstance(category_raw, str) or not category_raw.strip(): skipped += 1 continue category = category_raw.strip() if requested_categories is not None and category not in requested_categories: continue if args.limit is not None and samples_per_category.get(category, 0) >= args.limit: continue try: turns = extract_turns(row) except ValueError: skipped += 1 continue question_id = row.get("question_id") if not isinstance(question_id, str) or not question_id.strip(): skipped += 1 continue sample_id = question_id.strip() samples.append(Sample(id=sample_id, category=category, turns=turns)) samples_per_category[category] = samples_per_category.get(category, 0) + 1 if not samples: raise RuntimeError(f"no samples selected from bench={args.bench} category={args.category}") if skipped: print(f"speed_bench: skipped {skipped} rows without usable turns") return samples def parse_completion_response(data: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any], str | None, str]: usage = data.get("usage") or {} timings = data.get("timings") or {} finish_reason = None content = "" choices = data.get("choices") if isinstance(choices, list) and choices and isinstance(choices[0], dict): choice = choices[0] finish_reason = choice.get("finish_reason") message = choice.get("message") if isinstance(message, dict) and isinstance(message.get("content"), str): content = message["content"] elif isinstance(choice.get("text"), str): content = choice["text"] return usage, timings, finish_reason, content def run_request( endpoint: str, model: str | None, messages: list[dict[str, str]], osl: int, extra_inputs: dict[str, Any], timeout: float, ) -> tuple[dict[str, Any], float]: payload: dict[str, Any] = { "messages": messages, "max_tokens": osl, "stream": False, } if model: payload["model"] = model payload.update(extra_inputs) payload["max_tokens"] = osl start = time.perf_counter() response = requests.post(endpoint, json=payload, timeout=timeout) latency_s = time.perf_counter() - start if response.status_code != 200: body = response.text[:500].replace("\n", "\\n") raise RuntimeError(f"HTTP {response.status_code}: {body}") return response.json(), latency_s def run_one( sample: Sample, endpoint: str, model: str | None, osl: int, extra_inputs: dict[str, Any], timeout: float, ) -> RequestResult: selected_turns = sample.turns messages: list[dict[str, str]] = [] total_latency_s = 0.0 prompt_tokens = 0 completion_tokens = 0 total_tokens = 0 draft_n = 0 draft_n_accepted = 0 prompt_ms = 0.0 predicted_ms = 0.0 prompt_per_second = None predicted_per_second = None finish_reason: str | None = None try: for turn in selected_turns: messages.append({"role": "user", "content": turn}) data, latency_s = run_request(endpoint, model, messages, osl, extra_inputs, timeout) total_latency_s += latency_s usage, timings, finish_reason, assistant_text = parse_completion_response(data) turn_prompt_tokens = int(usage.get("prompt_tokens") or timings.get("prompt_n") or 0) turn_completion_tokens_count = int(usage.get("completion_tokens") or timings.get("predicted_n") or 0) turn_total_tokens_count = int(usage.get("total_tokens") or (turn_prompt_tokens + turn_completion_tokens_count)) prompt_tokens += turn_prompt_tokens completion_tokens += turn_completion_tokens_count total_tokens += turn_total_tokens_count draft_n += int(timings.get("draft_n") or 0) draft_n_accepted += int(timings.get("draft_n_accepted") or 0) prompt_ms += float(timings.get("prompt_ms") or 0) predicted_ms += float(timings.get("predicted_ms") or 0) if len(selected_turns) == 1 and isinstance(timings.get("prompt_per_second"), (int, float)): prompt_per_second = float(timings["prompt_per_second"]) if len(selected_turns) == 1 and isinstance(timings.get("predicted_per_second"), (int, float)): predicted_per_second = float(timings["predicted_per_second"]) messages.append({"role": "assistant", "content": assistant_text}) if total_tokens == 0: total_tokens = prompt_tokens + completion_tokens if len(selected_turns) > 1: prompt_per_second = (prompt_tokens / (prompt_ms / 1000)) if prompt_ms > 0 else None predicted_per_second = (completion_tokens / (predicted_ms / 1000)) if predicted_ms > 0 else None return RequestResult( id=sample.id, category=sample.category, ok=True, turns=len(selected_turns), latency_s=total_latency_s, prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=total_tokens, finish_reason=finish_reason, draft_n=draft_n, draft_n_accepted=draft_n_accepted, prompt_ms=prompt_ms if prompt_ms > 0 else None, predicted_ms=predicted_ms if predicted_ms > 0 else None, prompt_per_second=prompt_per_second, predicted_per_second=predicted_per_second, error=None, ) except Exception as exc: return RequestResult( id=sample.id, category=sample.category, ok=False, turns=len(selected_turns), latency_s=total_latency_s, prompt_tokens=0, completion_tokens=0, total_tokens=0, finish_reason=None, draft_n=0, draft_n_accepted=0, prompt_ms=None, predicted_ms=None, prompt_per_second=None, predicted_per_second=None, error=str(exc), ) def summarize_group(category: str, results: list[RequestResult]) -> dict[str, Any]: ok_results = [result for result in results if result.ok] latencies = [result.latency_s for result in ok_results] server_prompt_speeds = [ result.prompt_per_second for result in ok_results if result.prompt_per_second is not None ] server_completion_speeds = [ result.predicted_per_second for result in ok_results if result.predicted_per_second is not None ] turns = sum(result.turns for result in ok_results) draft_n = sum(result.draft_n for result in ok_results) accepted = sum(result.draft_n_accepted for result in ok_results) return { "category": category, "requests": len(ok_results), "turns": turns, "failed": len(results) - len(ok_results), "avg_prompt_t_s": statistics.mean(server_prompt_speeds) if server_prompt_speeds else None, "avg_pred_t_s": statistics.mean(server_completion_speeds) if server_completion_speeds else None, "avg_latency": statistics.mean(latencies) if latencies else None, "draft_n": draft_n, "accepted": accepted, "accept_rate": (accepted / draft_n) if draft_n > 0 else None, } def fmt_value(value: Any, kind: str = "") -> str: if value is None: return "n/a" if kind == "int": return str(int(value)) if kind == "rate": return f"{float(value):.4f}" if kind == "seconds": return f"{float(value):.3f}s" if kind == "speed": return f"{float(value):.2f}" if kind == "speedup": return f"{float(value):.2f}x" return str(value) def print_table(rows: list[dict[str, Any]]) -> None: columns = [ ("category", "category", ""), ("samples", "requests", "int"), ("avg_prompt_t/s", "avg_prompt_t_s", "speed"), ("avg_pred_t/s", "avg_pred_t_s", "speed"), ("avg_latency", "avg_latency", "seconds"), ("accept_rate", "accept_rate", "rate"), ] print_rows(rows, columns) def print_rows(rows: list[dict[str, Any]], columns: list[tuple[str, str, str]]) -> None: rendered_rows = [] for row in rows: rendered_rows.append([fmt_value(row.get(key), kind) for _, key, kind in columns]) widths = [len(header) for header, _, _ in columns] for rendered in rendered_rows: for i, cell in enumerate(rendered): widths[i] = max(widths[i], len(cell)) header = " ".join(header.ljust(widths[i]) for i, (header, _, _) in enumerate(columns)) print(header) print(" ".join("-" * width for width in widths)) for rendered in rendered_rows: print(" ".join(cell.ljust(widths[i]) for i, cell in enumerate(rendered))) def save_output(path: str, args: argparse.Namespace, samples: list[Sample], results: list[RequestResult], summary: list[dict[str, Any]]) -> None: payload = { "config": { "url": args.url, "model": args.model, "bench": args.bench, "category": args.category, "osl": args.osl, "concurrency": args.concurrency, "extra_inputs": args.extra_inputs, }, "selected_samples": len(samples), "completed_samples": sum(1 for result in results if result.ok), "failed_samples": sum(1 for result in results if not result.ok), "summary": summary, "results": [asdict(result) for result in results], } with open(path, "w", encoding="utf-8") as f: json.dump(payload, f, indent=2, sort_keys=True) def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser(description="Run SPEED-Bench against an OpenAI-compatible llama-server.") parser.add_argument("--url", default="localhost:8080", help="Server URL, for example localhost:8080 or http://localhost:8080/v1") parser.add_argument("--model", default=None, help="Optional model name to send in OpenAI requests") parser.add_argument("--bench", default="qualitative", help="SPEED-Bench config to run, for example qualitative or throughput_1k") parser.add_argument("--category", default="all", help="Category to run within the selected bench; use all for no category filter") parser.add_argument("--osl", type=int, default=4096, help="Output sequence length, mapped to max_tokens") parser.add_argument("--extra-inputs", default='{"temperature":0}', help="Extra request fields as a JSON object") parser.add_argument("--concurrency", type=int, default=1, help="Concurrent client requests; usually match llama-server --np") parser.add_argument("--limit", type=int, default=None, help="Optional sample limit per category for smoke tests") parser.add_argument("--timeout", type=float, default=600, help="Per-request timeout in seconds") parser.add_argument("--output", default=None, help="Optional path to save raw results JSON") args = parser.parse_args(argv) try: base_url = normalize_base_url(args.url) endpoint = base_url + "/chat/completions" extra_inputs = parse_extra_inputs(args.extra_inputs) args.extra_inputs = extra_inputs samples = load_samples(args) except Exception as exc: print(f"speed_bench: setup failed: {exc}", file=sys.stderr) return 2 print(f"speed_bench: loaded {len(samples)} samples from bench={args.bench} category={args.category}") results: list[RequestResult] = [] started = time.perf_counter() with concurrent.futures.ThreadPoolExecutor(max_workers=args.concurrency) as executor: futures = [ executor.submit(run_one, sample, endpoint, args.model, args.osl, extra_inputs, args.timeout) for sample in samples ] for future in tqdm(concurrent.futures.as_completed(futures), total=len(futures), desc="speed_bench", unit="sample"): result = future.result() results.append(result) elapsed = time.perf_counter() - started categories = list(dict.fromkeys(sample.category for sample in samples)) summary = [ summarize_group(category, [result for result in results if result.category == category]) for category in categories ] summary.append(summarize_group("overall", results)) print() print(f"Summary (elapsed={elapsed:.2f}s)") print_table(summary) if args.output: save_output(args.output, args, samples, results, summary) print(f"\nspeed_bench: wrote {args.output}") failed = sum(1 for result in results if not result.ok) if failed: print(f"\nspeed_bench: {failed} samples failed", file=sys.stderr) first_error = next((result.error for result in results if result.error), None) if first_error: print(f"first error: {first_error}", file=sys.stderr) return 1 return 0 if __name__ == "__main__": raise SystemExit(main())