#!/usr/bin/env python3
"""Measure prefill throughput as a function of prompt length.

The 128k plan lives or dies on this curve: at a flat 97 t/s a 128k prompt is
22 minutes, and on a hybrid model only part of the layers see the quadratic
term, so the shape is not obvious in advance.

Usage:
    python3 pp_probe.py --ctx 131072 --sizes 4096,8192,16384,32768 [-- extra flags]
"""
import argparse
import glob
import json
import os
import subprocess
import sys
import time
import http.client

BIN = os.path.expanduser("~/llama/llama-b11960/llama-server")
MODEL = os.path.expanduser("~/models/Qwen3.6-35B-A3B-MTP-UD-IQ3_XXS.gguf")
PORT = 8299

# Dense, repetitive but token-diverse filler: repeated verbatim so the
# tokenizer cost per token stays stable across sizes.
FILLER = (
    "The deployment pipeline builds the image, runs the unit tests, tags the "
    "result and pushes it to the registry. A failing stage stops the run and "
    "notifies the on-call engineer with the log tail attached. Retries are "
    "bounded, idempotent and observable, and every step records its duration. "
)


def post(path, payload, timeout=7200):
    c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=timeout)
    c.request("POST", path, body=json.dumps(payload),
              headers={"Content-Type": "application/json"})
    r = c.getresponse()
    body = r.read()
    c.close()
    if r.status != 200:
        raise RuntimeError(f"{path} {r.status}: {body[:300]!r}")
    return json.loads(body)


def mem_mb():
    vram = None
    for p in glob.glob("/sys/class/drm/card*/device/mem_info_vram_used"):
        try:
            vram = int(open(p).read()) // (2 ** 20)
        except (OSError, ValueError):
            pass
    return vram


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ctx", type=int, default=131072)
    ap.add_argument("--sizes", default="4096,8192,16384,32768")
    ap.add_argument("--model", default=MODEL)
    ap.add_argument("--threads", type=int, default=4)
    ap.add_argument("--ub", type=int, default=0, help="-ub (0 = default)")
    ap.add_argument("--b", type=int, default=0, help="-b (0 = default)")
    ap.add_argument("--np", type=int, default=1)
    ap.add_argument("--extra", default="")
    args = ap.parse_args()

    flags = ["-c", str(args.ctx), "-ngl", "99", "--n-cpu-moe", "999",
             "-fa", "on", "-ctk", "q8_0", "-ctv", "q8_0",
             "-t", str(args.threads), "-np", str(args.np), "--no-webui",
             "--spec-type", "draft-mtp", "--spec-draft-n-max", "1"]
    if args.ub:
        flags += ["-ub", str(args.ub)]
    if args.b:
        flags += ["-b", str(args.b)]
    if args.extra:
        flags += args.extra.split()

    log_path = os.path.expanduser("~/ngram/smoke/pp_probe.log")
    log = open(log_path, "wb")
    proc = subprocess.Popen([BIN, "-m", args.model, "--port", str(PORT),
                             *flags], stdout=log, stderr=subprocess.STDOUT,
                            stdin=subprocess.DEVNULL)
    print(f"flags: {' '.join(flags)}", file=sys.stderr)
    try:
        t0 = time.time()
        for _ in range(1200):
            if proc.poll() is not None:
                raise SystemExit(f"server died, see {log_path}")
            try:
                c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=5)
                c.request("GET", "/health")
                ok = c.getresponse().status == 200
                c.close()
                if ok:
                    break
            except OSError:
                pass
            time.sleep(1)
        print(f"server ready in {time.time()-t0:.1f}s, vram {mem_mb()} MB",
              file=sys.stderr)

        print(f"{'target':>8} {'actual':>8} {'pp t/s':>8} {'sec':>8} {'vram MB':>8}")
        for size in [int(s) for s in args.sizes.split(",")]:
            # ~3.6 chars per token for this filler; checked against /tokenize.
            text = FILLER * (size * 4 // len(FILLER) + 1)
            ntok = post("/tokenize", {"content": text,
                                      "add_special": False})["tokens"]
            # Trim to the target by words, then trust /tokenize over estimates.
            ids = ntok[:size]
            # llama-server accepts a token list only through /completion's
            # `prompt` when given as array of ids.
            out = post("/completion", {
                "prompt": ids, "n_predict": 1, "temperature": 0,
                "cache_prompt": False}, 7200)
            tm = out["timings"]
            print(f"{size:>8} {tm['prompt_n']:>8} "
                  f"{tm['prompt_n'] / tm['prompt_ms'] * 1000:>8.1f} "
                  f"{tm['prompt_ms']/1000:>8.1f} {mem_mb():>8}")
            sys.stdout.flush()
    finally:
        proc.terminate()
        try:
            proc.wait(timeout=60)
        except subprocess.TimeoutExpired:
            proc.kill()
        log.close()
        # The allocated KV/RS sizes are the other half of the 128k answer.
        with open(log_path, errors="replace") as f:
            for line in f:
                if "KV buffer size" in line or "RS buffer size" in line:
                    print(line.rstrip(), file=sys.stderr)


if __name__ == "__main__":
    main()