#!/usr/bin/env python3
"""
Real-corpus prompt-processing benchmark for llama-server on the RX 580 pod.

Runs ON THE POD. Normally invoked by bench.sh, not directly.

Starts llama-server from a given build directory with the production config,
waits for /health, POSTs fixed slices of the Pan Tadeusz corpus to /completion,
records timings.prompt_n / prompt_ms / prompt_per_second, then shuts the server
down cleanly.

Prompt slices are FIXED CONSTANTS (see SLICES below), calibrated once against
the Qwen3.6-35B-A3B tokenizer, so that every build is measured on byte-identical
input. Do not re-calibrate for a normal run; use --calibrate only if the model
or the corpus changes.

Usage:
  ppbench.py <build_dir> <label> [--sizes 4096,16384,32768] [--reps 3]
             [--results PATH] [--calibrate]
"""

import json
import os
import signal
import subprocess
import sys
import time
import urllib.request
from statistics import median

BENCH_DIR = "/root/bench"
CORPUS = os.path.join(BENCH_DIR, "pan-tadeusz.txt")
DEFAULT_RESULTS = os.path.join(BENCH_DIR, "results.txt")

MODEL = ("/root/.cache/huggingface/hub/models--unsloth--Qwen3.6-35B-A3B-GGUF/"
         "snapshots/a483e9e6cbd595906af30beda3187c2663a1118c/"
         "Qwen3.6-35B-A3B-UD-Q4_K_XL.gguf")

PORT = 8099
BASE = "http://127.0.0.1:%d" % PORT

# Production serving config. Keep in sync with /root/config.yaml (llama-swap).
# Traps encoded here on purpose:
#   -b / -ub 2048 : a ubatch larger than the prompt never fills, so every
#                   prompt length below must stay a multiple of 2048.
#   --no-mmap     : deprecated upstream in favour of --load-mode / -lm none,
#                   but still accepted; it is what production passes today.
SERVER_ARGS = ["-t", "6", "-ngl", "99", "--n-cpu-moe", "40",
               "-b", "2048", "-ub", "2048", "-fa", "1", "--no-mmap",
               "--ctx-size", "40960", "--no-warmup",
               "--port", str(PORT), "-m", MODEL]

# Calibrated once on 2026-09-08 against Qwen3.6-35B-A3B-UD-Q4_K_XL:
# Polish text of pan-tadeusz.txt runs 2.6702 chars/token, and /completion
# reports prompt_n identical to /tokenize (BOS offset 0).
# Each slice is corpus[0:NCHARS] read as UTF-8 text; NBYTES is the resulting
# UTF-8 byte length, recorded so the slice can be reproduced with byte tools.
# These produce EXACTLY the target prompt_n. Do not edit without recalibrating.
SLICES = {
    4096:  {"chars": 10776, "bytes": 11608},
    16384: {"chars": 43858, "bytes": 47175},
    32768: {"chars": 87165, "bytes": 93843},
}

DEFAULT_SIZES = [4096, 16384, 32768]


def post(path, payload, timeout=2400):
    req = urllib.request.Request(
        BASE + path,
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST")
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read().decode("utf-8"))


def health_ok():
    try:
        with urllib.request.urlopen(BASE + "/health", timeout=5) as r:
            return json.loads(r.read().decode("utf-8")).get("status") == "ok"
    except Exception:
        return False


def start_server(build_dir, logpath):
    env = dict(os.environ)
    env["LD_LIBRARY_PATH"] = os.path.join(build_dir, "bin")
    binary = os.path.join(build_dir, "bin", "llama-server")
    if not os.path.exists(binary):
        raise RuntimeError("no llama-server at %s" % binary)
    log = open(logpath, "wb")
    p = subprocess.Popen([binary] + SERVER_ARGS, stdout=log,
                         stderr=subprocess.STDOUT, env=env,
                         start_new_session=True)
    deadline = time.time() + 1200
    while time.time() < deadline:
        if p.poll() is not None:
            raise RuntimeError("llama-server exited early rc=%s, see %s"
                               % (p.returncode, logpath))
        if health_ok():
            return p
        time.sleep(2)
    stop_server(p)
    raise RuntimeError("llama-server not healthy within 1200s, see %s" % logpath)


def stop_server(p):
    if p is None or p.poll() is not None:
        return
    try:
        os.killpg(os.getpgid(p.pid), signal.SIGTERM)
    except Exception:
        p.terminate()
    for _ in range(120):
        if p.poll() is not None:
            return
        time.sleep(1)
    try:
        os.killpg(os.getpgid(p.pid), signal.SIGKILL)
    except Exception:
        p.kill()
    p.wait()


def n_tokens(text):
    """Exact token count via /tokenize. Cheap: no prefill."""
    return len(post("/tokenize", {"content": text}, timeout=300)["tokens"])


def calibrate(corpus, targets):
    """Recompute the SLICES table. Only needed if model or corpus changes."""
    probe = corpus[:2000]
    tk = n_tokens(probe)
    r = post("/completion", {"prompt": probe, "n_predict": 1,
                             "cache_prompt": False, "temperature": 0})
    offset = r["timings"]["prompt_n"] - tk
    cpt = len(probe) / float(tk)
    sys.stderr.write("calibrate: %.4f chars/token, prompt_n offset %d\n"
                     % (cpt, offset))
    for target in targets:
        want = target - offset
        lo, hi = 1, len(corpus)
        guess = min(len(corpus), max(1, int(want * cpt)))
        best = None
        for _ in range(60):
            got = n_tokens(corpus[:guess])
            if got == want:
                best = guess
                break
            if got < want:
                lo = guess + 1
            else:
                hi = guess - 1
            if lo > hi:
                break
            guess = (lo + hi) // 2
        if best is None:
            best = guess
        nbytes = len(corpus[:best].encode("utf-8"))
        sys.stderr.write("calibrate: %d tokens -> chars=%d bytes=%d\n"
                         % (target, best, nbytes))


def main():
    if len(sys.argv) < 3:
        sys.stderr.write(__doc__)
        return 2
    build_dir = os.path.abspath(sys.argv[1])
    label = sys.argv[2]
    sizes = list(DEFAULT_SIZES)
    reps = 3
    results_path = DEFAULT_RESULTS
    do_calibrate = "--calibrate" in sys.argv
    args = sys.argv[3:]
    for i, a in enumerate(args):
        if a == "--sizes":
            sizes = [int(x) for x in args[i + 1].split(",")]
        elif a == "--reps":
            reps = int(args[i + 1])
        elif a == "--results":
            results_path = args[i + 1]

    corpus = open(CORPUS, encoding="utf-8").read()
    logpath = os.path.join(BENCH_DIR, "server-%s-%d.log" % (label, int(time.time())))

    p = None
    try:
        sys.stderr.write("[%s] starting llama-server from %s\n" % (label, build_dir))
        t0 = time.time()
        p = start_server(build_dir, logpath)
        sys.stderr.write("[%s] healthy after %.1fs\n" % (label, time.time() - t0))

        if do_calibrate:
            calibrate(corpus, sizes)
            return 0

        stamp = time.strftime("%Y-%m-%dT%H:%M:%S")
        out = open(results_path, "a")
        for target in sizes:
            if target not in SLICES:
                sys.stderr.write("no calibrated slice for %d tokens, skipping\n" % target)
                continue
            nchars = SLICES[target]["chars"]
            prompt = corpus[:nchars]
            rates, mss, ns = [], [], set()
            for rep in range(1, reps + 1):
                r = post("/completion", {"prompt": prompt, "n_predict": 1,
                                         "cache_prompt": False, "temperature": 0})
                t = r["timings"]
                rates.append(t["prompt_per_second"])
                mss.append(t["prompt_ms"])
                ns.add(t["prompt_n"])
                line = ("RESULT corpus label=%s build=%s target=%d chars=%d "
                        "prompt_n=%d rep=%d prompt_ms=%.2f tps=%.2f time=%s"
                        % (label, build_dir, target, nchars, t["prompt_n"], rep,
                           t["prompt_ms"], t["prompt_per_second"], stamp))
                out.write(line + "\n")
                out.flush()
                sys.stderr.write(line + "\n")
            if ns != {target}:
                sys.stderr.write("WARNING: prompt_n %s != target %d; slice table "
                                 "is stale, rerun with --calibrate\n" % (sorted(ns), target))
            s = ("SUMMARY corpus label=%s target=%d prompt_n=%s median_tps=%.2f "
                 "min_tps=%.2f max_tps=%.2f median_ms=%.1f time=%s"
                 % (label, target, sorted(ns), median(rates), min(rates),
                    max(rates), median(mss), stamp))
            out.write(s + "\n")
            out.flush()
            sys.stderr.write(s + "\n")
        out.close()
    finally:
        stop_server(p)
        sys.stderr.write("[%s] server stopped\n" % label)
    return 0


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