#!/usr/bin/env python3
"""Measure prefill, decode and memory for one OpenAI-compatible server.

Purpose: send the same two prompts to whichever engine is listening, time the first token
    and the tokens after it from the streaming response, derive prefill from the difference
    between a short and a long prompt, read what the machine reports about memory while the
    model is loaded (as a difference from an idle snapshot taken before any server started),
    and append one JSON line per test to the lab notebook so that engines on the same machine
    can be compared honestly.
Platform: all (Python standard library only; the memory readers use /proc/meminfo, nvidia-smi,
    rocm-smi or vm_stat, whichever exist)
Minimum memory: 12 GB
Assumes: Python 3.9 or later; a server already running and reachable at --base-url, serving
    the model named by --model; for memory deltas, an idle snapshot written earlier with
    --snapshot-only while no model server was running.

Usage: python3 compare-engines.py --snapshot-only idle-memory.json
       python3 compare-engines.py --engine llama.cpp --base-url http://127.0.0.1:8080/v1 \
           --model Qwen3-8B-Q4_K_M --quant Q4_K_M --context-length 8192 \
           --engine-version "0.4.0 (build 12345)" --backend CUDA --host "DGX Spark" \
           --server-pid "$(cat llama-server.pid)" --idle-snapshot idle-memory.json \
           --labbook labbook.md

Method, in the order the script applies it:
  * Two tests. "short" is a one-line question; "long" repeats a fixed paragraph so that the
    prompt is some 1,500 tokens longer. Both ask for --max-tokens tokens.
  * Every request starts with a random run tag. Engines keep the KV cache of recent prompts
    and skip prefill for a prefix they have seen (llama-server, vLLM, TensorRT-LLM and
    mlx_lm.server all do this by default), so an identical prompt sent twice measures the
    cache, not prefill. A different first line makes every prompt new. The server's own
    count of reused tokens is recorded as max_cached_prompt_tokens so that you can check.
  * One unmeasured warm-up request per test, then --repetitions measured ones; medians.
  * Time to first token (ttft) is from sending the request to the first streamed chunk that
    carries text: answer text, reasoning text (reasoning_content or reasoning) or a tool call.
    Qwen3 thinks before answering, and servers put that text in different fields.
  * Decode rate is (completion tokens - 1) / (time from first to last text chunk), because
    prefill produces the first token. Completion tokens come from the server's usage block
    when it sends one, and from counting chunks otherwise (token_source says which).
  * Prefill rate is taken from the difference between the tests:
        (prompt_tokens_long - prompt_tokens_short) / (ttft_long - ttft_short)
    which cancels the fixed costs (HTTP, scheduling, the first decode step) that a single
    ttft includes. llama-server also reports its own rates in a "timings" block; when present
    they are recorded as server_prefill_tokens_per_s and server_decode_tokens_per_s.
"""

from __future__ import annotations

import argparse
import json
import platform
import shutil
import statistics
import subprocess
import sys
import time
import urllib.error
import urllib.request
import uuid
from pathlib import Path

LAB = "part-08/lab-same-model-every-engine"

SHORT_PROMPT = "In two sentences, say what a KV cache is and why it grows with context length."

LONG_PARAGRAPH = (
    "A language model reads its prompt in one pass and then writes its answer one token at a "
    "time. The first phase is compute bound because every weight that is read is reused across "
    "many tokens of the prompt. The second phase is bandwidth bound because every weight has to "
    "be read again for each single token that is produced. Any measurement that reports one "
    "number for both phases has thrown away the distinction that explains how the machine "
    "behaves. "
)

FOLLOW_UP = "Summarise the passage above in exactly three sentences."

TEXT_FIELDS = ("content", "reasoning_content", "reasoning")


def long_prompt(repeats: int) -> str:
    """A deterministic prompt of `repeats` copies of the paragraph, ending in an instruction."""
    return (LONG_PARAGRAPH * repeats) + "\n\n" + FOLLOW_UP


def headers(api_key: str, accept: str) -> dict:
    out = {"Content-Type": "application/json", "Accept": accept}
    if api_key:
        out["Authorization"] = f"Bearer {api_key}"
    return out


def wait_for_server(args) -> None:
    """Poll GET {base}/models until it answers 200, so a slow load is not measured as a failure."""
    url = args.base_url.rstrip("/") + "/models"
    deadline = time.monotonic() + args.wait_seconds
    last = "no answer yet"
    while True:
        request = urllib.request.Request(url, headers=headers(args.api_key, "application/json"))
        try:
            with urllib.request.urlopen(request, timeout=10) as response:  # noqa: S310 - local server
                if response.status == 200:
                    return
                last = f"status {response.status}"
        except urllib.error.HTTPError as exc:
            last = f"status {exc.code}"
        except (urllib.error.URLError, TimeoutError, OSError) as exc:
            last = str(exc)
        if time.monotonic() >= deadline:
            raise RuntimeError(f"{url} did not answer 200 within {args.wait_seconds} s ({last})")
        time.sleep(2)


def post_stream(base_url: str, api_key: str, body: dict, timeout: float):
    """POST a streaming chat completion and yield (arrival_time, parsed_chunk) pairs."""
    url = base_url.rstrip("/") + "/chat/completions"
    data = json.dumps(body).encode("utf-8")
    request = urllib.request.Request(
        url, data=data, headers=headers(api_key, "text/event-stream"), method="POST"
    )
    with urllib.request.urlopen(request, timeout=timeout) as response:  # noqa: S310 - local server
        while True:
            raw = response.readline()
            if not raw:
                return
            line = raw.decode("utf-8", "replace").strip()
            if not line.startswith("data:"):
                continue
            payload = line[len("data:"):].strip()
            if payload == "[DONE]":
                return
            try:
                yield time.perf_counter(), json.loads(payload)
            except json.JSONDecodeError:
                continue


def one_run(args, prompt: str) -> dict:
    """One streamed completion, timed. Returns raw timings and counts rather than rates."""
    content = prompt if args.reuse_prompt else f"[run {uuid.uuid4().hex[:12]}]\n{prompt}"
    body = {
        "model": args.model,
        "messages": [{"role": "user", "content": content}],
        "max_tokens": args.max_tokens,
        "temperature": 0,
        "stream": True,
        "stream_options": {"include_usage": True},
    }
    started = time.perf_counter()
    first_at = None
    last_at = None
    chunks = 0
    usage = None
    timings = None
    finish_reason = None
    fields_seen = set()

    for arrived, chunk in post_stream(args.base_url, args.api_key, body, args.timeout):
        if chunk.get("usage"):
            usage = chunk["usage"]
        if chunk.get("timings"):
            timings = chunk["timings"]
        for choice in chunk.get("choices") or []:
            delta = choice.get("delta") or {}
            carried = [f for f in TEXT_FIELDS if isinstance(delta.get(f), str) and delta.get(f)]
            if delta.get("tool_calls"):
                carried.append("tool_calls")
            if carried:
                fields_seen.update(carried)
                chunks += 1
                if first_at is None:
                    first_at = arrived
                last_at = arrived
            if choice.get("finish_reason"):
                finish_reason = choice["finish_reason"]

    if first_at is None:
        raise RuntimeError("the server streamed no text; check the model name and the server log")

    usage = usage or {}
    details = usage.get("prompt_tokens_details") or {}
    cached = details.get("cached_tokens")
    if cached is None and timings:
        cached = timings.get("cache_n")
    return {
        "ttft_s": first_at - started,
        "decode_s": max(last_at - first_at, 1e-9),
        "chunks": chunks,
        "prompt_tokens": usage.get("prompt_tokens"),
        "completion_tokens": usage.get("completion_tokens"),
        "cached_tokens": cached,
        "finish_reason": finish_reason,
        "fields": sorted(fields_seen),
        "server_prefill": (timings or {}).get("prompt_per_second"),
        "server_decode": (timings or {}).get("predicted_per_second"),
    }


def rates(run: dict) -> dict:
    """Turn one run's timings into rates, saying where the token counts came from."""
    if run["completion_tokens"]:
        produced, source = run["completion_tokens"], "usage"
    else:
        produced, source = run["chunks"], "chunks"
    after_first = max(produced - 1, 1)
    return {
        "ttft_s": run["ttft_s"],
        "decode_tokens_per_s": after_first / run["decode_s"],
        "prompt_tokens": run["prompt_tokens"],
        "completion_tokens": produced,
        "token_source": source,
        "cached_tokens": run["cached_tokens"],
        "finish_reason": run["finish_reason"],
        "fields": run["fields"],
        "server_prefill": run["server_prefill"],
        "server_decode": run["server_decode"],
    }


def run_command(command: list) -> str | None:
    """Run a short informational command; return its output, or None if it is unavailable."""
    if not shutil.which(command[0]):
        return None
    try:
        out = subprocess.run(command, capture_output=True, text=True, timeout=20, check=False)
    except (OSError, subprocess.SubprocessError):
        return None
    return out.stdout.strip() if out.returncode == 0 else None


def to_float(text) -> float | None:
    try:
        return float(str(text).strip())
    except (TypeError, ValueError):
        return None


def memory_snapshot() -> dict:
    """Every memory figure this machine will report, in MB (1 MB = 1,000,000 bytes)."""
    snap = {"taken_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())}

    meminfo = Path("/proc/meminfo")
    if meminfo.exists():
        values = {}
        for line in meminfo.read_text(encoding="utf-8").splitlines():
            key, _, rest = line.partition(":")
            parts = rest.split()
            if parts and parts[0].isdigit():
                values[key] = int(parts[0]) * 1024 / 1e6  # the file reports kB (KiB)
        snap["mem_total_mb"] = values.get("MemTotal")
        snap["mem_available_mb"] = values.get("MemAvailable")
        snap["swap_free_mb"] = values.get("SwapFree")

    text = run_command(["nvidia-smi", "--query-gpu=memory.used,memory.total",
                        "--format=csv,noheader,nounits"])
    if text:
        fields = [f.strip() for f in text.splitlines()[0].split(",")]
        snap["gpu_used_mb"] = to_float(fields[0]) if fields else None  # MiB; null on a GB10
        snap["gpu_total_mb"] = to_float(fields[1]) if len(fields) > 1 else None
        # Per-process figures are recorded only when at least one row parses to a number. Under
        # WSL2, NVIDIA's CUDA on WSL guide lists "active compute process" among the NVML
        # queries not yet supported, so the query returns no usable rows there; recording 0
        # would make every server look as if it added nothing.
        apps = run_command(["nvidia-smi", "--query-compute-apps=pid,used_memory",
                            "--format=csv,noheader,nounits"])
        if apps is not None:
            used = [to_float(line.split(",")[1]) for line in apps.splitlines() if "," in line]
            used = [u for u in used if u is not None]
            if used:
                snap["gpu_process_mb"] = sum(used)

    text = run_command(["rocm-smi", "--showmeminfo", "vram", "gtt", "--json"])
    if text:
        try:
            cards = json.loads(text)
        except json.JSONDecodeError:
            cards = {}
        vram = gtt = 0.0
        for card in cards.values() if isinstance(cards, dict) else []:
            vram += to_float(card.get("VRAM Total Used Memory (B)")) or 0.0
            gtt += to_float(card.get("GTT Total Used Memory (B)")) or 0.0
        snap["rocm_vram_used_mb"] = vram / 1e6
        snap["rocm_gtt_used_mb"] = gtt / 1e6

    text = run_command(["vm_stat"])
    if text:
        page = 4096
        pages = {}
        for line in text.splitlines():
            if "page size of" in line:
                digits = [t for t in line.replace("(", " ").split() if t.isdigit()]
                page = int(digits[0]) if digits else page
            key, _, rest = line.partition(":")
            rest = rest.strip().rstrip(".")
            if rest.isdigit():
                pages[key.strip()] = int(rest)
        mb = lambda k: pages.get(k, 0) * page / 1e6  # noqa: E731 - tiny local helper
        snap["mac_free_mb"] = mb("Pages free")
        snap["mac_wired_mb"] = mb("Pages wired down")
        snap["mac_compressed_mb"] = mb("Pages occupied by compressor")
        snap["mac_used_mb"] = mb("Pages active") + snap["mac_wired_mb"] + snap["mac_compressed_mb"]
        swap = run_command(["sysctl", "-n", "vm.swapusage"])
        if swap and "used =" in swap:
            snap["mac_swap_used_mb"] = to_float(swap.split("used =")[1].split()[0].rstrip("M"))
    return snap


def memory_delta(idle: dict, loaded: dict) -> dict:
    """Memory the server added, per source: loaded minus idle (idle minus loaded for 'available')."""
    out = {}
    if idle.get("mem_available_mb") is not None and loaded.get("mem_available_mb") is not None:
        out["mem_available_mb"] = round(idle["mem_available_mb"] - loaded["mem_available_mb"])
    for key in ("gpu_used_mb", "gpu_process_mb", "rocm_vram_used_mb", "rocm_gtt_used_mb",
                "mac_wired_mb", "mac_used_mb"):
        before = idle.get(key)
        if key == "gpu_process_mb" and before is None and idle.get("gpu_total_mb"):
            before = 0.0  # nvidia-smi answered at idle and listed no compute process
        if before is not None and loaded.get(key) is not None:
            out[key] = round(loaded[key] - before)
    return out


def process_resident_mb(pid: int) -> float | None:
    """Resident set size of the server process from ps, in MB. Only the process itself."""
    if not pid:
        return None
    text = run_command(["ps", "-o", "rss=", "-p", str(pid)])
    try:
        return round(float(text.split()[0]) * 1024 / 1e6) if text else None
    except (IndexError, ValueError):
        return None


def median(runs: list, key: str):
    values = [r[key] for r in runs if r.get(key) is not None]
    return statistics.median(values) if values else None


def measure(args, label: str, prompt: str) -> dict:
    """Warm up once, repeat the test, take medians."""
    if not args.no_warmup:
        one_run(args, prompt)
    runs = [rates(one_run(args, prompt)) for _ in range(args.repetitions)]
    cached = [r["cached_tokens"] for r in runs if r["cached_tokens"] is not None]
    return {
        "lab": LAB,
        "engine": args.engine,
        "engine_version": args.engine_version,
        "backend": args.backend,
        "host": args.host,
        "model": args.model,
        "quant": args.quant,
        "context_length": args.context_length,
        "test": label,
        "repetitions": args.repetitions,
        "max_tokens": args.max_tokens,
        "ttft_s": round(median(runs, "ttft_s"), 4),
        "decode_tokens_per_s": round(median(runs, "decode_tokens_per_s"), 2),
        "prompt_tokens": median(runs, "prompt_tokens"),
        "completion_tokens": median(runs, "completion_tokens"),
        "token_source": runs[-1]["token_source"],
        "finish_reason": runs[-1]["finish_reason"],
        "text_fields": runs[-1]["fields"],
        "max_cached_prompt_tokens": max(cached) if cached else None,
        "prompt_tag": not args.reuse_prompt,
        "server_prefill_tokens_per_s": median(runs, "server_prefill"),
        "server_decode_tokens_per_s": median(runs, "server_decode"),
        "measured_on": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
    }


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--snapshot-only", metavar="FILE",
                        help="write a memory snapshot to FILE and exit (run with no server up)")
    parser.add_argument("--engine", help="engine name as it should appear in the notebook")
    parser.add_argument("--base-url", default="http://127.0.0.1:8080/v1", help="OpenAI-compatible base URL")
    parser.add_argument("--model", help="model name to send in every request")
    parser.add_argument("--api-key", default="", help="bearer token, when the server wants one")
    parser.add_argument("--engine-version", default="", help="the version string you read from the engine")
    parser.add_argument("--backend", default="", help="CUDA, Metal, Vulkan, ROCm, and so on")
    parser.add_argument("--host", default="", help="short description of the machine")
    parser.add_argument("--quant", default="", help="quantisation of the weights being served")
    parser.add_argument("--context-length", type=int, default=0, help="context the server was started with")
    parser.add_argument("--server-pid", type=int, default=0, help="process id of the server, for ps")
    parser.add_argument("--idle-snapshot", default="", help="snapshot file written by --snapshot-only")
    parser.add_argument("--max-tokens", type=int, default=128, help="tokens to generate per run")
    parser.add_argument("--long-repeats", type=int, default=16, help="paragraph repeats in the long prompt")
    parser.add_argument("--repetitions", type=int, default=3, help="measured runs per test")
    parser.add_argument("--timeout", type=float, default=600.0, help="seconds to wait for one response")
    parser.add_argument("--wait-seconds", type=int, default=120, help="seconds to wait for the server to answer")
    parser.add_argument("--no-warmup", action="store_true", help="skip the unmeasured first run")
    parser.add_argument("--reuse-prompt", action="store_true",
                        help="omit the run tag, so repeats hit the prefix cache (a demonstration, not a measurement)")
    parser.add_argument("--labbook", default="labbook.md", help="notebook to append to")
    parser.add_argument("--print-only", action="store_true", help="print the results, record nothing")
    args = parser.parse_args()

    if args.snapshot_only:
        snap = memory_snapshot()
        snap["system"] = f"{platform.system()} {platform.machine()}"
        Path(args.snapshot_only).write_text(json.dumps(snap, indent=2) + "\n", encoding="utf-8")
        readings = {k: round(v) for k, v in snap.items() if isinstance(v, float)}
        print(f"    idle snapshot written to {args.snapshot_only}")
        print(f"    {json.dumps(readings)}")
        return 0

    if not args.engine or not args.model:
        parser.error("--engine and --model are required unless --snapshot-only is given")

    idle = {}
    if args.idle_snapshot:
        try:
            idle = json.loads(Path(args.idle_snapshot).read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError) as exc:
            print(f"    cannot read {args.idle_snapshot}: {exc}", file=sys.stderr)
            return 1

    tests = [("short", SHORT_PROMPT), ("long", long_prompt(args.long_repeats))]
    records = []
    try:
        wait_for_server(args)
        for label, prompt in tests:
            records.append(measure(args, label, prompt))
    except (urllib.error.URLError, RuntimeError, TimeoutError, OSError) as exc:
        print(f"    FAILED: {exc}", file=sys.stderr)
        print("    Usually: the server is not up yet, the model name differs from GET /v1/models,",
              file=sys.stderr)
        print("    or the server needs --api-key.", file=sys.stderr)
        return 1

    loaded = memory_snapshot()
    delta = memory_delta(idle, loaded) if idle else {}
    rss = process_resident_mb(args.server_pid)
    short, long_ = records
    marginal = None
    if (long_["prompt_tokens"] and short["prompt_tokens"]
            and long_["ttft_s"] > short["ttft_s"]):
        marginal = round((long_["prompt_tokens"] - short["prompt_tokens"])
                         / (long_["ttft_s"] - short["ttft_s"]), 1)
    for record in records:
        record["prefill_tokens_per_s"] = marginal
        record["memory"] = loaded
        record["memory_delta_mb"] = delta
        record["process_resident_mb"] = rss

    for r in records:
        prompt = f"{r['prompt_tokens']:.0f}" if r["prompt_tokens"] is not None else "?"
        print(f"    {args.engine:>14} {r['test']:>5}  prompt {prompt:>5} tok  ttft {r['ttft_s']:.3f} s  "
              f"decode {r['decode_tokens_per_s']:.2f} tok/s  "
              f"{r['completion_tokens']:.0f} tok ({r['finish_reason']}), counts from {r['token_source']}, "
              f"cached {r['max_cached_prompt_tokens']}")
    if marginal is not None:
        extra = long_["prompt_tokens"] - short["prompt_tokens"]
        print(f"    {args.engine:>14} prefill from the difference: {marginal:.1f} tok/s over {extra:.0f} extra prompt tokens")
    else:
        print(f"    {args.engine:>14} prefill from the difference: not computable (no prompt token counts)")
    if long_["server_prefill_tokens_per_s"] is not None:
        print(f"    {args.engine:>14} server's own timings, long test: prefill "
              f"{long_['server_prefill_tokens_per_s']:.1f} tok/s, decode {long_['server_decode_tokens_per_s']:.2f} tok/s")
    print(f"    {args.engine:>14} memory added since the idle snapshot (MB): {json.dumps(delta) if delta else 'no idle snapshot given'}")
    print(f"    {args.engine:>14} server process resident set (MB): {rss}")

    if args.print_only:
        return 0
    notebook = Path(args.labbook)
    if not notebook.exists():
        print(f"    {notebook} does not exist; creating it", file=sys.stderr)
    with notebook.open("a", encoding="utf-8") as handle:
        for record in records:
            handle.write(json.dumps(record) + "\n")
    print(f"    recorded {len(records)} line(s) in {notebook}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())