#!/usr/bin/env python3
"""Common LLM-serving load generator — drives any OpenAI-compatible endpoint
(vLLM or SGLang) identically, so the vLLM-vs-SGLang comparison is apples-to-apples.

Streams responses to measure TTFT (time-to-first-token) and TPOT (time-per-output-
token), not just aggregate throughput. Run it against each engine in turn on the
SAME exclusive GPU, same prompt profile, same concurrency sweep.

Usage:
  python3 bench.py --url http://localhost:8006/v1 --model granite-4.1-8b-nvfp4 \
      --concurrency 1 10 50 100 --in-tokens 2048 --out-tokens 256 --requests 200

For the prefill-heavy agent-memory regime, bump --in-tokens to ~30000.
"""
import argparse, json, time, urllib.request, statistics, concurrent.futures

def stream_one(url, model, prompt, out_tokens):
    body = json.dumps({"model": model, "messages": [{"role": "user", "content": prompt}],
                       "max_tokens": out_tokens, "temperature": 0.7, "stream": True}).encode()
    req = urllib.request.Request(url + "/chat/completions", data=body,
                                 headers={"Content-Type": "application/json"})
    t0 = time.perf_counter(); ttft = None; n = 0
    with urllib.request.urlopen(req, timeout=600) as r:
        for line in r:
            line = line.strip()
            if not line or not line.startswith(b"data:"):
                continue
            if line == b"data: [DONE]":
                break
            try:
                delta = json.loads(line[5:])["choices"][0]["delta"].get("content")
            except Exception:
                delta = None
            if delta:
                if ttft is None:
                    ttft = time.perf_counter() - t0
                n += 1
    dur = time.perf_counter() - t0
    return ttft, n, dur

def run(url, model, prompt, out_tokens, N):
    t0 = time.perf_counter()
    with concurrent.futures.ThreadPoolExecutor(max_workers=N) as ex:
        res = list(ex.map(lambda _: stream_one(url, model, prompt, out_tokens), range(N)))
    wall = time.perf_counter() - t0
    res = [r for r in res if r[1] > 0]
    ttfts = sorted(r[0] for r in res if r[0])
    tot_tok = sum(r[1] for r in res)
    # TPOT: per-request (dur - ttft) / (tokens - 1), averaged
    tpots = [(d - t) / max(1, n - 1) for (t, n, d) in res if t and n > 1]
    pc = lambda xs, p: xs[min(len(xs) - 1, int(len(xs) * p))] if xs else float("nan")
    return {
        "N": N, "ok": len(res),
        "agg_tok_s": tot_tok / wall,
        "ttft_p50": pc(ttfts, 0.50), "ttft_p99": pc(ttfts, 0.99),
        "tpot_ms": (statistics.mean(tpots) * 1000) if tpots else float("nan"),
    }

if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--url", required=True)
    ap.add_argument("--model", required=True)
    ap.add_argument("--concurrency", type=int, nargs="+", default=[1, 10, 50, 100])
    ap.add_argument("--in-tokens", type=int, default=2048)
    ap.add_argument("--out-tokens", type=int, default=256)
    a = ap.parse_args()
    # crude prompt sized to ~in-tokens (word ≈ 1.3 tokens)
    prompt = "Summarize and analyze the following text in detail. " + ("data point " * int(a.in_tokens / 1.3))
    print(f"# {a.url}  model={a.model}  in≈{a.in_tokens} out={a.out_tokens}")
    print(f"{'N':>4} {'ok':>4} {'agg_tok/s':>10} {'ttft_p50':>9} {'ttft_p99':>9} {'tpot_ms':>8}")
    for N in a.concurrency:
        r = run(a.url, a.model, prompt, a.out_tokens, N)
        print(f"{r['N']:>4} {r['ok']:>4} {r['agg_tok_s']:>10.0f} {r['ttft_p50']:>9.2f} {r['ttft_p99']:>9.2f} {r['tpot_ms']:>8.1f}")