#!/usr/bin/env python3
# bench.py - llama-server benchmark: prompt processing / generation speed / VRAM (llama.cpp HTTP API).
# Usage:     python3 bench.py <port> <label> <out.json>      e.g.  python3 bench.py 8081 q8_0 vk_res_q8_0.json
# Needs:     python3 + requests, rocm-smi (VRAM readout only), a llama-server already listening on 127.0.0.1:<port>.
# Protocol:  prompts of 1k, 4k, 8k, 16k and 30k tokens; 1 warm-up + 7 passes each; seed 42, temperature 0;
#            256 tokens generated per pass (ignore_eos); prompt cache disabled; /no_think. Median [min-max] is computed afterwards.
# Server:    --ctx-size 32768 --n-gpu-layers 99 --flash-attn on (see run_vk.sh / vulkan_command.txt).
# Tested on: Ubuntu 24.04, AMD RX 9070 XT 16 GB, Mesa 25.2.8 RADV, llama.cpp b11177 (commit 1ab7e5ad2,
#            official ggml-org archive), Qwen3 14B Q4_K_M.
"""Bench llama-server : lecture du prompt / génération / VRAM. Usage: bench.py <port> <label> <out.json>"""
import json, random, re, statistics as st, subprocess, sys, requests

PORT, LABEL, OUT = int(sys.argv[1]), sys.argv[2], sys.argv[3]
URL = f"http://127.0.0.1:{PORT}"
SIZES = [1000, 4000, 8000, 16000, 30000]
WARMUP, RUNS, SEED, NGEN = 1, 7, 42, 256

def tok(text):
    return requests.post(f"{URL}/tokenize", json={"content": text, "add_special": False}).json()["tokens"]
def detok(t):
    return requests.post(f"{URL}/detokenize", json={"tokens": t}).json()["content"]

def vram():
    o = subprocess.run(["rocm-smi", "--showmeminfo", "vram", "-d", "0"], capture_output=True, text=True).stdout
    tot = int(re.search(r"Total Memory \(B\):\s*(\d+)", o).group(1))
    used = int(re.search(r"Total Used Memory \(B\):\s*(\d+)", o).group(1))
    return used, tot - used

# Corpus pseudo-aléatoire déterministe (mots anglais courants, graine fixe)
rnd = random.Random(1234)
vocab = ("time year people way day man thing woman life child world school state family student group country "
         "problem hand part place case week company system program question work government number night point "
         "home water room mother area money story fact month lot right study book eye job word business issue "
         "side kind head house service friend father power hour game line end member law car city community").split()
corpus = " ".join(rnd.choice(vocab) for _ in range(40000))
ctoks = tok(corpus)
PRE = "<|im_start|>user\nBelow is a list of random words. Count nothing, just write a short story.\n\n"
POST = "\n\n/no_think<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
overhead = len(tok(PRE + POST))

res = {"label": LABEL, "seed": SEED, "n_gen": NGEN, "runs": RUNS, "warmup": WARMUP, "sizes": {}}
for n in SIZES:
    body = detok(ctoks[: n - overhead])
    prompt = PRE + body + POST
    real = len(tok(prompt))
    pp, tg, ng, npr = [], [], [], []
    vr = []
    for i in range(WARMUP + RUNS):
        r = requests.post(f"{URL}/completion", json={
            "prompt": prompt, "n_predict": NGEN, "seed": SEED, "temperature": 0.0,
            "cache_prompt": False, "ignore_eos": True, "stream": False}, timeout=900).json()
        t = r["timings"]
        u, f = vram()
        print(f"{LABEL} n={n} (réel {real}) run {i} {'[chauffe]' if i < WARMUP else ''} "
              f"pp={t['prompt_per_second']:.1f} tg={t['predicted_per_second']:.2f} "
              f"prompt_n={t['prompt_n']} pred_n={t['predicted_n']} vram_used={u/2**30:.2f}GiB free={f/2**30:.2f}GiB", flush=True)
        if i >= WARMUP:
            pp.append(t["prompt_per_second"]); tg.append(t["predicted_per_second"])
            ng.append(t["predicted_n"]); npr.append(t["prompt_n"]); vr.append((u, f))
    res["sizes"][n] = {"prompt_tokens": real, "prompt_n_processed": npr, "gen_n": ng, "pp": pp, "tg": tg, "vram": vr}
json.dump(res, open(OUT, "w"), indent=1)