120B benchmark script: tokens/s, TTFT, full-spread (gx10)

benchmark_120b_tokens_per_second.py

#!/usr/bin/env python3
"""
Benchmark tokens per second and TTFT for the gx10 120B LLM endpoint (Qwen 122B).

Calls POST /v1/chat/completions (non-streaming for throughput; optional streaming
for TTFT). Supports multiple runs (mean ± std) and optional concurrent requests.

Usage:
  python3 scripts/benchmark_120b_tokens_per_second.py [BASE_URL]
  python3 scripts/benchmark_120b_tokens_per_second.py --runs 5 --ttft --concurrent 2
  python3 scripts/benchmark_120b_tokens_per_second.py --full-spread   # TTFT + variance + C=1,2,4
  GX10_120B_URL=http://localhost:8002 python3 scripts/benchmark_120b_tokens_per_second.py

Defaults: BASE_URL = http://gx10-83fb.tail3dac72.ts.net:8002
"""

from __future__ import annotations

import argparse
import json
import os
import statistics
import sys
import threading
import time
import urllib.request
from datetime import datetime, timezone

DEFAULT_BASE_URL = "http://gx10-83fb.tail3dac72.ts.net:8002"
BENCHMARK_PROMPT = (
    "Count from 1 to 50, one number per line. Do not add any other text or explanation."
)
MAX_TOKENS = 256


def get_model_id(base_url: str) -> str:
    req = urllib.request.Request(
        f"{base_url.rstrip('/')}/v1/models",
        headers={"Accept": "application/json"},
    )
    with urllib.request.urlopen(req, timeout=30) as resp:
        data = json.load(resp)
    models = data.get("data", data.get("models", []))
    if not models:
        raise SystemExit("No models returned from /v1/models")
    first = models[0]
    return first.get("id", first.get("model", first.get("name", "unknown")))


def run_completion(
    base_url: str,
    model_id: str,
    *,
    stream: bool = False,
) -> tuple[dict | None, float]:
    """
    Run one completion. If stream=True, returns (None, ttft_seconds) and does not
    parse full response. If stream=False, returns (response_json, wall_seconds).
    """
    url = f"{base_url.rstrip('/')}/v1/chat/completions"
    body = {
        "model": model_id,
        "messages": [{"role": "user", "content": BENCHMARK_PROMPT}],
        "max_tokens": MAX_TOKENS,
        "stream": stream,
    }
    data = json.dumps(body).encode("utf-8")
    headers = {
        "Content-Type": "application/json",
        "Accept": "text/event-stream" if stream else "application/json",
    }
    req = urllib.request.Request(url, data=data, method="POST", headers=headers)
    start = time.perf_counter()
    with urllib.request.urlopen(req, timeout=300) as resp:
        if not stream:
            raw = resp.read()
            elapsed = time.perf_counter() - start
            return json.loads(raw.decode("utf-8")), elapsed
        # Streaming: read until first content chunk for TTFT
        ttft = None
        buf = b""
        while True:
            chunk = resp.read(8192)
            if not chunk:
                break
            buf += chunk
            while b"\n\n" in buf and ttft is None:
                event, _, buf = buf.partition(b"\n\n")
                event = event.strip()
                if not event or event == b"data: [DONE]":
                    continue
                if event.startswith(b"data: "):
                    try:
                        j = json.loads(event[6:].decode("utf-8"))
                        for c in j.get("choices", []):
                            if c.get("delta", {}).get("content"):
                                ttft = time.perf_counter() - start
                                break
                    except (json.JSONDecodeError, KeyError):
                        pass
            if ttft is not None:
                break
        elapsed = time.perf_counter() - start
        return (None, ttft if ttft is not None else elapsed)


def extract_usage(result: dict) -> tuple[int, int]:
    usage = result.get("usage") or {}
    prompt_tokens = usage.get("prompt_tokens", 0)
    completion_tokens = usage.get("completion_tokens")
    if completion_tokens is None:
        choices = result.get("choices", [])
        if choices:
            msg = choices[0].get("message", {})
            content = msg.get("content") or msg.get("reasoning_content") or ""
            completion_tokens = max(1, len(content) // 4)
        else:
            completion_tokens = 0
    return prompt_tokens, completion_tokens


def run_one_throughput(base_url: str, model_id: str) -> tuple[float, float, int]:
    """Returns (elapsed_sec, tokens_per_sec, completion_tokens)."""
    result, elapsed = run_completion(base_url, model_id, stream=False)
    if result is None:
        raise RuntimeError("Unexpected streaming response in throughput run")
    _, completion_tokens = extract_usage(result)
    tps = completion_tokens / elapsed if elapsed > 0 else 0.0
    return elapsed, tps, completion_tokens


def measure_ttft(base_url: str, model_id: str) -> float | None:
    """One streaming request; returns TTFT in seconds or None if streaming unsupported."""
    try:
        _, ttft_or_elapsed = run_completion(base_url, model_id, stream=True)
        return ttft_or_elapsed
    except Exception:
        return None


def run_concurrent(
    base_url: str,
    model_id: str,
    concurrency: int,
) -> tuple[list[tuple[float, float, int]], float]:
    """Launch concurrency requests in parallel. Returns (per_request_results, wall_elapsed)."""
    results: list[tuple[float, float, int]] = []
    errors: list[Exception] = []
    start = time.perf_counter()

    def one() -> None:
        try:
            r = run_one_throughput(base_url, model_id)
            results.append(r)
        except Exception as e:
            errors.append(e)

    threads = [threading.Thread(target=one) for _ in range(concurrency)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()
    wall = time.perf_counter() - start
    if errors:
        raise errors[0]
    return results, wall


def _run_full_spread(base_url: str, model_id: str, quiet: bool, record_period_path: str | None) -> None:
    """Run all permutations: warmup, TTFT, then C=1, C=2, C=4 each with runs=5."""
    out = sys.stderr if not quiet else sys.stdout
    runs = 5
    results: list[tuple[str, float, float | None, str]] = []  # (label, mean_tps, stdev_or_ttft, extra)
    start_wall = time.perf_counter()
    start_iso = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")

    def section(title: str) -> None:
        if not quiet:
            print(file=out)
            print(f"=== {title} ===", file=out)

    # 1. Warmup
    section("Warmup")
    if not quiet:
        print("One completion (no timing)...", file=out)
    try:
        run_one_throughput(base_url, model_id)
    except Exception as e:
        print(f"Warmup failed: {e}", file=out)
        sys.exit(1)

    # 2. TTFT
    section("TTFT (time to first token)")
    if not quiet:
        print("One streaming request...", file=out)
    ttft = measure_ttft(base_url, model_id)
    if ttft is not None:
        results.append(("TTFT (s)", 0.0, ttft, f"{ttft:.2f} s"))
        if not quiet:
            print(f"  TTFT: {ttft:.2f} s", file=out)
    else:
        if not quiet:
            print("  TTFT: (streaming not supported or failed)", file=out)

    # 3. C=1, runs=5 (variance)
    section(f"Throughput C=1 (runs={runs})")
    tps_list: list[float] = []
    for i in range(runs):
        elapsed, tps, n = run_one_throughput(base_url, model_id)
        tps_list.append(tps)
    mean_tps = statistics.mean(tps_list)
    stdev_tps = statistics.stdev(tps_list) if len(tps_list) > 1 else 0.0
    results.append((f"C=1 (t/s)", mean_tps, stdev_tps, f"{mean_tps:.2f} ± {stdev_tps:.2f}"))
    if not quiet:
        print(f"  Output t/s: {mean_tps:.2f} ± {stdev_tps:.2f}", file=out)

    # 4. C=2, runs=5
    section(f"Throughput C=2 (runs={runs})")
    agg_tps_list: list[float] = []
    for _ in range(runs):
        per_req, wall = run_concurrent(base_url, model_id, 2)
        total_tokens = sum(r[2] for r in per_req)
        agg_tps_list.append(total_tokens / wall if wall > 0 else 0.0)
    mean_agg = statistics.mean(agg_tps_list)
    stdev_agg = statistics.stdev(agg_tps_list) if len(agg_tps_list) > 1 else 0.0
    results.append((f"C=2 (t/s)", mean_agg, stdev_agg, f"{mean_agg:.2f} ± {stdev_agg:.2f}"))
    if not quiet:
        print(f"  Aggregate t/s: {mean_agg:.2f} ± {stdev_agg:.2f}", file=out)

    # 5. C=4, runs=5
    section(f"Throughput C=4 (runs={runs})")
    agg_tps_list = []
    for _ in range(runs):
        per_req, wall = run_concurrent(base_url, model_id, 4)
        total_tokens = sum(r[2] for r in per_req)
        agg_tps_list.append(total_tokens / wall if wall > 0 else 0.0)
    mean_agg = statistics.mean(agg_tps_list)
    stdev_agg = statistics.stdev(agg_tps_list) if len(agg_tps_list) > 1 else 0.0
    results.append((f"C=4 (t/s)", mean_agg, stdev_agg, f"{mean_agg:.2f} ± {stdev_agg:.2f}"))
    if not quiet:
        print(f"  Aggregate t/s: {mean_agg:.2f} ± {stdev_agg:.2f}", file=out)

    # Summary table
    section("Full-spread summary")
    if not quiet:
        print(f"  {'Metric':<14} {'Result':<20}", file=out)
        for label, _m, _s, extra in results:
            print(f"  {label:<14} {extra}", file=out)
    # One-line summary for quiet / scripting
    if quiet:
        for label, _m, _s, extra in results:
            print(f"{label}\t{extra}")

    end_iso = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
    if record_period_path:
        os.makedirs(os.path.dirname(record_period_path) or ".", exist_ok=True)
        period_data: dict = {"start_iso": start_iso, "end_iso": end_iso}
        # Include results for charts/reports (ttft_s, throughput by concurrency)
        period_data["benchmark"] = {}
        for label, mean_val, std_or_ttft, _extra in results:
            if "TTFT" in label:
                period_data["benchmark"]["ttft_s"] = std_or_ttft
            elif label == "C=1 (t/s)":
                period_data["benchmark"]["c1_tps_mean"] = mean_val
                period_data["benchmark"]["c1_tps_std"] = std_or_ttft or 0.0
            elif label == "C=2 (t/s)":
                period_data["benchmark"]["c2_tps_mean"] = mean_val
                period_data["benchmark"]["c2_tps_std"] = std_or_ttft or 0.0
            elif label == "C=4 (t/s)":
                period_data["benchmark"]["c4_tps_mean"] = mean_val
                period_data["benchmark"]["c4_tps_std"] = std_or_ttft or 0.0
        with open(record_period_path, "w", encoding="utf-8") as f:
            json.dump(period_data, f, indent=2)
        if not quiet:
            print(file=out)
            print(f"Benchmark period: {start_iso} → {end_iso}", file=out)
            print(f"  Saved to {record_period_path}", file=out)
            print("  Correlate: python3 scripts/correlate_benchmark_prometheus.py --period-file", record_period_path, file=out)
            print("  Charts:    python3 scripts/benchmark_prometheus_report.py --period-file", record_period_path, "--output docs/evidence/benchmark_report.html", file=out)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Benchmark output tokens/s and optional TTFT for gx10 120B endpoint."
    )
    parser.add_argument(
        "base_url",
        nargs="?",
        default=os.environ.get("GX10_120B_URL", DEFAULT_BASE_URL),
        help="Base URL of the 120B API (default: gx10 Tailscale:8002)",
    )
    parser.add_argument(
        "-q",
        "--quiet",
        action="store_true",
        help="Only print summary (tokens/s, duration, tokens)",
    )
    parser.add_argument(
        "-n",
        "--runs",
        type=int,
        default=1,
        metavar="N",
        help="Number of throughput runs for mean ± std (default: 1)",
    )
    parser.add_argument(
        "--warmup",
        action="store_true",
        help="Do one warmup completion before timed runs",
    )
    parser.add_argument(
        "--ttft",
        action="store_true",
        help="Measure time-to-first-token with one streaming request",
    )
    parser.add_argument(
        "-c",
        "--concurrent",
        type=int,
        default=1,
        metavar="C",
        help="Run C requests in parallel (default: 1); reports aggregate throughput",
    )
    parser.add_argument(
        "--full-spread",
        action="store_true",
        help="Run all permutations: warmup, TTFT, then C=1/2/4 each with --runs 5 --warmup",
    )
    parser.add_argument(
        "--record-period",
        metavar="PATH",
        help="With --full-spread: write start_iso/end_iso to JSON file for correlate_benchmark_prometheus.py",
    )
    args = parser.parse_args()
    base_url = args.base_url.rstrip("/")

    if not args.quiet:
        print(f"Base URL: {base_url}", file=sys.stderr)
    try:
        model_id = get_model_id(base_url)
    except Exception as e:
        print(f"Failed to get model id: {e}", file=sys.stderr)
        sys.exit(1)
    if not args.quiet:
        print(f"Model:   {model_id}", file=sys.stderr)

    # Full-spread: all permutations (warmup, TTFT, C=1/2/4 with runs=5)
    if args.full_spread:
        record_path = args.record_period
        if record_path and not os.path.isabs(record_path):
            repo = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
            record_path = os.path.join(repo, record_path)
        _run_full_spread(base_url, model_id, args.quiet, record_path or None)
        return

    # Optional TTFT (streaming)
    if args.ttft:
        if not args.quiet:
            print("Measuring TTFT (streaming)...", file=sys.stderr)
        ttft = measure_ttft(base_url, model_id)
        if ttft is not None:
            if not args.quiet:
                print(f"  TTFT: {ttft:.2f} s", file=sys.stderr)
            else:
                print(f"ttft_s\t{ttft:.2f}")
        else:
            if not args.quiet:
                print("  TTFT: (streaming not supported or failed)", file=sys.stderr)

    # Warmup
    if args.warmup and not args.quiet:
        print("Warmup run...", file=sys.stderr)
    if args.warmup:
        try:
            run_one_throughput(base_url, model_id)
        except Exception as e:
            print(f"Warmup failed: {e}", file=sys.stderr)
            sys.exit(1)

    # Throughput: concurrent with multiple rounds, or single concurrent round
    if args.concurrent > 1:
        if args.runs > 1:
            # N rounds of C concurrent; report mean ± std of aggregate t/s
            if not args.quiet:
                print(f"Concurrent C={args.concurrent}, runs={args.runs}...", file=sys.stderr)
            agg_tps_list: list[float] = []
            for round_num in range(args.runs):
                try:
                    per_req, wall = run_concurrent(base_url, model_id, args.concurrent)
                    total_tokens = sum(r[2] for r in per_req)
                    agg_tps = total_tokens / wall if wall > 0 else 0.0
                    agg_tps_list.append(agg_tps)
                except Exception as e:
                    print(f"Round {round_num + 1} failed: {e}", file=sys.stderr)
                    sys.exit(1)
            mean_tps = statistics.mean(agg_tps_list)
            stdev_tps = statistics.stdev(agg_tps_list) if len(agg_tps_list) > 1 else 0.0
            if args.quiet:
                print(f"{mean_tps:.2f}\t±{stdev_tps:.2f}\t(concurrent {args.concurrent}, {args.runs} rounds)")
            else:
                print(file=sys.stderr)
                print(f"Concurrent C={args.concurrent} ({args.runs} rounds):", file=sys.stderr)
                print(f"  Aggregate t/s: {mean_tps:.2f} ± {stdev_tps:.2f}", file=sys.stderr)
                print(file=sys.stderr)
                print(f"Aggregate output tokens/s: {mean_tps:.2f} ± {stdev_tps:.2f}")
            return
        # Single round of C concurrent
        if not args.quiet:
            print(f"Concurrent runs (C={args.concurrent})...", file=sys.stderr)
        try:
            per_req, wall = run_concurrent(base_url, model_id, args.concurrent)
        except Exception as e:
            print(f"Concurrent run failed: {e}", file=sys.stderr)
            sys.exit(1)
        total_tokens = sum(r[2] for r in per_req)
        agg_tps = total_tokens / wall if wall > 0 else 0.0
        if args.quiet:
            print(f"{agg_tps:.2f}\t{wall:.2f}s\t{total_tokens} tokens (concurrent {args.concurrent})")
        else:
            print(file=sys.stderr)
            print("Concurrent results:", file=sys.stderr)
            print(f"  Requests:     {len(per_req)}", file=sys.stderr)
            print(f"  Total tokens: {total_tokens}", file=sys.stderr)
            print(f"  Wall time:    {wall:.2f} s", file=sys.stderr)
            print(f"  Aggregate t/s: {agg_tps:.2f}", file=sys.stderr)
            for i, (elapsed, tps, n) in enumerate(per_req, 1):
                print(f"  Request {i}: {tps:.2f} t/s, {elapsed:.2f} s, {n} tokens", file=sys.stderr)
            print(file=sys.stderr)
            print(f"Aggregate output tokens/s: {agg_tps:.2f}")
        return

    # Single-threaded runs
    runs_tps: list[float] = []
    runs_elapsed: list[float] = []
    runs_tokens: list[int] = []
    for i in range(args.runs):
        try:
            elapsed, tps, n = run_one_throughput(base_url, model_id)
            runs_elapsed.append(elapsed)
            runs_tps.append(tps)
            runs_tokens.append(n)
        except Exception as e:
            print(f"Run {i + 1} failed: {e}", file=sys.stderr)
            sys.exit(1)

    if args.quiet:
        if args.runs == 1:
            print(f"{runs_tps[0]:.2f}\t{runs_elapsed[0]:.2f}s\t{runs_tokens[0]} tokens")
        else:
            mean_tps = statistics.mean(runs_tps)
            print(f"{mean_tps:.2f}\t{statistics.mean(runs_elapsed):.2f}s\t{sum(runs_tokens)} tokens ({args.runs} runs)")
        return

    # Report
    print(file=sys.stderr)
    print("Results:", file=sys.stderr)
    print(f"  Runs:           {args.runs}", file=sys.stderr)
    print(f"  Prompt:         {BENCHMARK_PROMPT[:50]}...", file=sys.stderr)
    total_completion = sum(runs_tokens)
    print(f"  Completion tokens (total): {total_completion}", file=sys.stderr)
    mean_elapsed = statistics.mean(runs_elapsed)
    print(f"  Wall time (mean): {mean_elapsed:.2f} s", file=sys.stderr)
    mean_tps = statistics.mean(runs_tps)
    print(f"  Output tokens/s (mean): {mean_tps:.2f}", file=sys.stderr)
    if args.runs > 1:
        print(f"  Output tokens/s (std):  {statistics.stdev(runs_tps):.2f}", file=sys.stderr)
        print(f"  Wall time (std):       {statistics.stdev(runs_elapsed):.2f} s", file=sys.stderr)
    print(file=sys.stderr)
    print(f"Output tokens per second: {mean_tps:.2f}")


if __name__ == "__main__":
    main()