#!/usr/bin/env python3
"""Cache-busted CONCURRENT throughput bench against a vLLM OpenAI seat.

Companion to quickbench.py, which is bs=1 only. Single-stream tok/s and
aggregate throughput are different questions: one is latency-bound, the other is
batch-bound, and a seat can win on one while losing the other.

MEASUREMENT TRAPS THIS AVOIDS (all three have produced confident wrong results
here before — see docs/pfi/model-quantization-playbook.md §5):

  1. Prefix caching fakes speed. Every request gets a FRESH UNSEEDED nonce, so
     no prompt is ever repeated. Never seed the cache-buster: a seeded nonce
     regenerates the previous run's prompts verbatim and reads cache-hit
     throughput (~41k tok/s) instead of real prefill (~5k).
  2. MTP acceptance is read from /metrics as a DELTA across the run, not as a
     cumulative total, or a long-lived seat's history swamps the measurement.
  3. Aggregate throughput is computed from the WALL CLOCK of the whole batch,
     not the sum of per-request rates — the latter double-counts overlap and
     reports a number the seat cannot actually deliver.

Usage:
    uv run concbench.py --base http://10.250.50.54:8015 \
        --model qwen3.8-27b-uncensored --concurrency 1 --concurrency 6
"""
import argparse, json, random, statistics, string, sys, time
import urllib.request
from concurrent.futures import ThreadPoolExecutor

TOPICS = [
    "a coral reef ecosystem", "the Roman aqueduct system",
    "how lithium-ion cells degrade", "the history of the printing press",
    "how radar altimeters work", "glacier mass balance",
    "the design of the Saturn V F-1 engine", "how sourdough fermentation works",
    "the metallurgy of Damascus steel", "how noise-cancelling headphones work",
    "the ecology of mangrove forests", "the development of the marine chronometer",
]


def nonce(n=14):
    # Unseeded on purpose. A seeded RNG reproduces the previous run's prompts
    # and turns this into a prefix-cache benchmark.
    return "".join(random.choices(string.ascii_lowercase + string.digits, k=n))


def spec_counters(base):
    """(drafted, accepted) from /metrics, or (None, None) if unavailable."""
    try:
        with urllib.request.urlopen(base + "/metrics", timeout=30) as r:
            text = r.read().decode()
    except Exception:
        return None, None
    d = a = None
    for line in text.splitlines():
        if line.startswith("#"):
            continue
        if line.startswith("vllm:spec_decode_num_draft_tokens_total"):
            d = float(line.rsplit(" ", 1)[1])
        elif line.startswith("vllm:spec_decode_num_accepted_tokens_total"):
            a = float(line.rsplit(" ", 1)[1])
    return d, a


def one(base, model, max_tokens, timeout):
    body = {
        "model": model,
        "messages": [{"role": "user", "content":
                      f"[req-{nonce()}] Explain {random.choice(TOPICS)} in detail."}],
        "max_tokens": max_tokens,
        "temperature": 0.7,
    }
    req = urllib.request.Request(
        base + "/v1/chat/completions", data=json.dumps(body).encode(),
        headers={"Content-Type": "application/json"})
    t0 = time.time()
    with urllib.request.urlopen(req, timeout=timeout) as r:
        d = json.load(r)
    dt = time.time() - t0
    return d["usage"]["completion_tokens"], dt


def run(base, model, conc, reqs, max_tokens, timeout):
    d0, a0 = spec_counters(base)
    t0 = time.time()
    with ThreadPoolExecutor(max_workers=conc) as ex:
        out = list(ex.map(lambda _: one(base, model, max_tokens, timeout), range(reqs)))
    wall = time.time() - t0
    d1, a1 = spec_counters(base)

    toks = sum(t for t, _ in out)
    lat = [dt for _, dt in out]
    acc = None
    if d0 is not None and a0 is not None and d1 is not None and a1 is not None:
        drafted = d1 - d0
        if drafted > 0:
            acc = (a1 - a0) / drafted * 100
    return {
        "concurrency": conc, "requests": reqs, "wall_s": round(wall, 2),
        "completion_tokens": toks,
        "aggregate_tok_s": round(toks / wall, 2),           # wall clock, not sum-of-rates
        "per_stream_tok_s": round(toks / wall / conc, 2),
        "median_latency_s": round(statistics.median(lat), 2),
        "mtp_acceptance_pct": round(acc, 1) if acc is not None else None,
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", default="http://10.250.50.54:8015")
    ap.add_argument("--model", default="qwen3.8-27b-uncensored")
    ap.add_argument("--concurrency", type=int, action="append", default=[])
    ap.add_argument("--requests-per-stream", type=int, default=4)
    ap.add_argument("--max-tokens", type=int, default=400)
    ap.add_argument("--timeout", type=int, default=900)
    ap.add_argument("--tag", default="run")
    ap.add_argument("--out", default=None)
    a = ap.parse_args()
    concs = a.concurrency or [1, 6]

    results = []
    for c in concs:
        reqs = c * a.requests_per_stream
        r = run(a.base, a.model, c, reqs, a.max_tokens, a.timeout)
        results.append(r)
        print(f"  conc={c:<3} reqs={reqs:<3} wall={r['wall_s']:>7.2f}s  "
              f"aggregate={r['aggregate_tok_s']:>8.2f} tok/s  "
              f"per-stream={r['per_stream_tok_s']:>7.2f}  "
              f"median-lat={r['median_latency_s']:>6.2f}s  "
              f"MTP={r['mtp_acceptance_pct']}%")

    if a.out:
        json.dump({"tag": a.tag, "model": a.model, "base": a.base,
                   "max_tokens": a.max_tokens, "results": results},
                  open(a.out, "w"), indent=2)
    return 0


if __name__ == "__main__":
    sys.exit(main())