#!/usr/bin/env python3
"""
Lab 01 — Quantize & Benchmark harness.

Measures TTFT, TPOT, throughput, and peak VRAM against an OpenAI-compatible
endpoint (vLLM / TGI / TensorRT-LLM-via-Triton all expose one). Run the same
script against an FP16, INT8, and INT4 server and tabulate the results.

Why an external endpoint instead of in-process generate()? Because production
performance is a *serving* property (continuous batching, scheduling), not a
single-call property — see Knowledge 04 and Knowledge 06 §8. Single-call timing
overstates latency and understates throughput.

Usage:
    # 1. start a server, e.g.:
    #    vllm serve meta-llama/Meta-Llama-3-8B-Instruct --port 8000
    #    vllm serve <model> --quantization fp8 --port 8000        # INT8/FP8 run
    #    vllm serve <awq-model> --quantization awq --port 8000     # INT4 run
    # 2. run this against it:
    python run_benchmark.py --base-url http://localhost:8000/v1 \
        --model meta-llama/Meta-Llama-3-8B-Instruct \
        --concurrency 32 --requests 256 --max-tokens 256 --label FP16

Dependencies: pip install openai numpy
VRAM: read separately via `nvidia-smi --query-gpu=memory.used --format=csv` while loaded.
"""
import argparse
import asyncio
import statistics
import time

from openai import AsyncOpenAI

PROMPT = (
    "Explain the roofline performance model and why LLM token generation is "
    "memory-bandwidth bound rather than compute bound. Be precise and concrete."
)


async def one_request(client, model, max_tokens):
    """Fire one streaming request; return (ttft, tpot, n_tokens)."""
    t0 = time.perf_counter()
    first_token_t = None
    n = 0
    stream = await client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": PROMPT}],
        max_tokens=max_tokens,
        temperature=0.0,  # fixed sampling so runs are comparable (K06 §7)
        stream=True,
    )
    async for chunk in stream:
        if not chunk.choices:
            continue
        delta = chunk.choices[0].delta.content
        if delta:
            now = time.perf_counter()
            if first_token_t is None:
                first_token_t = now            # TTFT boundary
            n += 1
    t_end = time.perf_counter()
    ttft = (first_token_t - t0) if first_token_t else float("nan")
    # TPOT = mean inter-token latency over the decode phase
    tpot = ((t_end - first_token_t) / max(n - 1, 1)) if first_token_t else float("nan")
    return ttft, tpot, n


async def run(args):
    client = AsyncOpenAI(base_url=args.base_url, api_key="not-needed")

    # Warmup (discard — CUDA init, autotune, cold caches; K06 §7 rule 1).
    await asyncio.gather(*[one_request(client, args.model, 16) for _ in range(args.concurrency)])

    sem = asyncio.Semaphore(args.concurrency)

    async def guarded():
        async with sem:
            return await one_request(client, args.model, args.max_tokens)

    wall0 = time.perf_counter()
    results = await asyncio.gather(*[guarded() for _ in range(args.requests)])
    wall = time.perf_counter() - wall0

    ttfts = [r[0] for r in results]
    tpots = [r[1] for r in results]
    total_tokens = sum(r[2] for r in results)

    def pct(xs, p):
        xs = sorted(x for x in xs if x == x)  # drop NaN
        if not xs:
            return float("nan")
        k = min(len(xs) - 1, int(round((p / 100) * (len(xs) - 1))))
        return xs[k]

    print(f"\n=== {args.label}  (concurrency={args.concurrency}, requests={args.requests}) ===")
    print(f"  p50 TTFT  : {pct(ttfts,50)*1000:8.1f} ms")
    print(f"  p95 TTFT  : {pct(ttfts,95)*1000:8.1f} ms")
    print(f"  p50 TPOT  : {pct(tpots,50)*1000:8.2f} ms   ({1/pct(tpots,50):6.1f} tok/s/stream)")
    print(f"  p95 TPOT  : {pct(tpots,95)*1000:8.2f} ms")
    print(f"  mean TPOT : {statistics.mean([t for t in tpots if t==t])*1000:8.2f} ms")
    print(f"  THROUGHPUT: {total_tokens/wall:8.1f} output tok/s (aggregate)")
    print(f"  wall      : {wall:8.2f} s  for {total_tokens} tokens")
    print("  (record peak VRAM separately via nvidia-smi while the server is loaded)")
    print("  --> add this row to the Lab 01 table; repeat for INT8/FP8 and INT4 servers.\n")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base-url", default="http://localhost:8000/v1")
    ap.add_argument("--model", required=True)
    ap.add_argument("--concurrency", type=int, default=32)
    ap.add_argument("--requests", type=int, default=256)
    ap.add_argument("--max-tokens", type=int, default=256)
    ap.add_argument("--label", default="FP16", help="precision label for the output row")
    asyncio.run(run(ap.parse_args()))


if __name__ == "__main__":
    main()