#!/usr/bin/env python3
"""
Decode benchmark for vLLM's OpenAI-compatible endpoint (allsmartllm.com engineering).
Used for the articles "Two models on one DGX Spark" and "Qwen3.8-27B on one DGX Spark".

Streams via SSE so that time to first token is measured, not estimated.
Reasoning and content tokens are counted separately: a client that only
counts content sees far too low a throughput when thinking is enabled.
Single stream: (completion tokens - 1) / (last chunk - first chunk).
Concurrent runs: total completion tokens / wall time.

Usage:
  ./bench.py --label baseline-fp8 --out ergebnis.json
  ./bench.py --port 8003 --label nvfp4 --concurrency 1,4
"""
import argparse, json, statistics, sys, time, urllib.request
from concurrent.futures import ThreadPoolExecutor

# Echte Aufgaben aus dem mid/vision/coder-Einsatz, keine Zufallstoken.
# Bewusst gemischt: kurze und lange Eingaben, Code und Prosa.
PROMPTS = [
    "Schreib eine Python-Funktion, die eine SSE-Antwort zeilenweise parst und "
    "reasoning- von content-Deltas trennt. Mit Docstring und Typannotationen.",

    "Wir betreiben vLLM auf zwei gekoppelten DGX Sparks mit Tensor-Parallelitaet "
    "ueber 200 GbE. Der Durchsatz liegt bei der Haelfte des Erwarteten. Nenne die "
    "fuenf wahrscheinlichsten Ursachen, nach Pruefaufwand sortiert.",

    "Erklaer den Unterschied zwischen FP8 e4m3 und e5m2 bei der Inferenz von "
    "grossen Sprachmodellen. Wann nimmt man welches Format?",

    "Refactor dieses Bash-Fragment, sodass es bei jedem Fehler abbricht und "
    "temporaere Dateien zuverlaessig aufraeumt:\n"
    "cp $1 /tmp/x && process /tmp/x > $2 && rm /tmp/x",

    "Ein systemd-Dienst startet einen Docker-Container mit --rm. Der Container "
    "stuerzt sporadisch ab und niemand weiss warum. Was ist das Problem und wie "
    "loest man es, ohne den Dienst neu zu bauen?",

    "Schreib eine SQL-Abfrage, die pro Kunde den letzten Auftrag der letzten "
    "90 Tage liefert, inklusive Kunden ohne Auftrag. Tabellen: kunden(id, name), "
    "auftraege(id, kunde_id, datum, betrag).",

    "Fasse zusammen, was Multi-Token-Prediction beim Decoding bringt und unter "
    "welchen Bedingungen es sich nicht lohnt. Sei konkret bei den Kosten.",

    "Uebersetz diese Fehlermeldung ins Deutsche und erklaer die Ursache in zwei "
    "Saetzen: 'ValueError: No available memory for the cache blocks. Try "
    "increasing gpu_memory_utilization when initializing the engine.'",

    "Schreib einen regulaeren Ausdruck, der IPv4-Adressen in Logzeilen findet, "
    "aber Versionsnummern wie 1.2.3.4 in Klammern ausschliesst. Erklaer ihn.",

    "Entwirf das Datenmodell fuer ein Benchmark-Archiv: mehrere Modelle, mehrere "
    "Quantisierungen, Messreihen mit Parallelitaetsgrad und Zeitstempel. Nenne "
    "Tabellen, Schluessel und die zwei wichtigsten Indizes.",
]


def one_request(url, model, prompt, max_tokens, thinking, timeout):
    """Eine Anfrage per Streaming. Gibt Zeiten und Tokenzahlen zurueck."""
    body = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": max_tokens,
        "temperature": 0,
        "stream": True,
        "stream_options": {"include_usage": True},
    }
    if not thinking:
        body["chat_template_kwargs"] = {"enable_thinking": False}

    req = urllib.request.Request(
        url, data=json.dumps(body).encode("utf-8"),
        headers={"Content-Type": "application/json"})

    t0 = time.perf_counter()
    first = first_content = last = None
    usage = None
    n_reason_chunks = n_content_chunks = 0

    with urllib.request.urlopen(req, timeout=timeout) as resp:
        for raw in resp:
            line = raw.decode("utf-8", "replace").strip()
            if not line.startswith("data:"):
                continue
            payload = line[5:].strip()
            if payload == "[DONE]":
                break
            d = json.loads(payload)
            now = time.perf_counter()
            if d.get("usage"):
                usage = d["usage"]
            choices = d.get("choices") or []
            if not choices:
                continue
            delta = choices[0].get("delta") or {}
            # vLLM nennt das Feld je nach Version reasoning oder reasoning_content
            r = delta.get("reasoning") or delta.get("reasoning_content")
            c = delta.get("content")
            if r:
                n_reason_chunks += 1
            if c:
                n_content_chunks += 1
            if r or c:
                if first is None:
                    first = now
                last = now
            if c and first_content is None:
                first_content = now

    if usage is None or first is None or last is None or last <= first:
        return None  # unbrauchbar, wird als nicht messbar gewertet

    comp = usage.get("completion_tokens", 0)
    details = usage.get("completion_tokens_details") or {}
    reason_tok = details.get("reasoning_tokens")

    return {
        "ttft_s": first - t0,
        "ttft_content_s": (first_content - t0) if first_content else None,
        "gen_s": last - first,
        "wall_s": last - t0,
        "prompt_tokens": usage.get("prompt_tokens", 0),
        "completion_tokens": comp,
        "reasoning_tokens": reason_tok,
        "reason_chunks": n_reason_chunks,
        "content_chunks": n_content_chunks,
        # -1: das erste Token gehoert zum TTFT, nicht zur Decode-Phase
        "decode_tps": (comp - 1) / (last - first) if comp > 1 else None,
    }


def run_batch(url, model, c, max_tokens, thinking, timeout):
    """c Anfragen parallel. Liefert Aggregat ueber die ganze Welle."""
    prompts = [PROMPTS[i % len(PROMPTS)] for i in range(c)]
    t0 = time.perf_counter()
    with ThreadPoolExecutor(max_workers=c) as ex:
        res = list(ex.map(
            lambda p: one_request(url, model, p, max_tokens, thinking, timeout),
            prompts))
    wall = time.perf_counter() - t0
    ok = [r for r in res if r]
    if not ok:
        return None
    total_comp = sum(r["completion_tokens"] for r in ok)
    return {
        "concurrency": c,
        "ok": len(ok),
        "failed": len(res) - len(ok),
        "wall_s": wall,
        "aggregate_tps": total_comp / wall,
        "per_stream_tps": statistics.median(
            [r["decode_tps"] for r in ok if r["decode_tps"]]),
        "ttft_median_s": statistics.median([r["ttft_s"] for r in ok]),
        "ttft_content_median_s": statistics.median(
            [r["ttft_content_s"] for r in ok if r["ttft_content_s"]]
        ) if any(r["ttft_content_s"] for r in ok) else None,
        "completion_tokens_total": total_comp,
        "reasoning_tokens_total": sum(
            r["reasoning_tokens"] or 0 for r in ok),
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--host", default="localhost")
    ap.add_argument("--port", type=int, default=8001)
    ap.add_argument("--model", default="qwen3.6-35b")
    ap.add_argument("--label", required=True)
    ap.add_argument("--concurrency", default="1,4,8,16,32")
    ap.add_argument("--runs", type=int, default=3)
    ap.add_argument("--max-tokens", type=int, default=512)
    ap.add_argument("--timeout", type=int, default=900)
    ap.add_argument("--no-thinking", action="store_true",
                    help="enable_thinking=false mitmessen")
    ap.add_argument("--out", default=None)
    args = ap.parse_args()

    url = f"http://{args.host}:{args.port}/v1/chat/completions"
    thinking = not args.no_thinking
    levels = [int(x) for x in args.concurrency.split(",")]

    print(f"# {args.label}  ({url}, model={args.model}, "
          f"thinking={'an' if thinking else 'aus'}, max_tokens={args.max_tokens})",
          flush=True)

    # Warmup, wird verworfen
    print("Warmup ...", flush=True)
    run_batch(url, args.model, 1, args.max_tokens, thinking, args.timeout)

    out = {"label": args.label, "url": url, "model": args.model,
           "thinking": thinking, "max_tokens": args.max_tokens,
           "runs": args.runs, "ts": time.strftime("%Y-%m-%d %H:%M:%S"),
           "levels": {}}

    for c in levels:
        runs = []
        for i in range(args.runs):
            r = run_batch(url, args.model, c, args.max_tokens, thinking,
                          args.timeout)
            if r:
                runs.append(r)
                print(f"  c={c:<3} lauf {i+1}: aggregat "
                      f"{r['aggregate_tps']:.1f} tok/s, pro strom "
                      f"{r['per_stream_tps']:.1f}, ttft "
                      f"{r['ttft_median_s']*1000:.0f} ms", flush=True)
            else:
                print(f"  c={c:<3} lauf {i+1}: NICHT MESSBAR", flush=True)
        if not runs:
            out["levels"][str(c)] = {"messbar": False}
            continue
        med = lambda k: statistics.median([r[k] for r in runs if r[k] is not None]) \
            if any(r[k] is not None for r in runs) else None
        out["levels"][str(c)] = {
            "messbar": True,
            "laeufe": len(runs),
            "aggregate_tps": med("aggregate_tps"),
            "per_stream_tps": med("per_stream_tps"),
            "ttft_s": med("ttft_median_s"),
            "ttft_content_s": med("ttft_content_median_s"),
            "completion_tokens_total": med("completion_tokens_total"),
            "reasoning_tokens_total": med("reasoning_tokens_total"),
            "failed": sum(r["failed"] for r in runs),
        }

    js = json.dumps(out, indent=2, ensure_ascii=False)
    if args.out:
        with open(args.out, "w") as f:
            f.write(js)
        print(f"\ngeschrieben nach {args.out}")
    else:
        print(js)


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