llama.cpp/tools/server/bench/speed-bench/speed_bench.py

#!/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())