serverAppConfigs/scripts/autoresearch/frank-primary-measure.py

#!/usr/bin/env python3
"""frank-primary decode measurement — runs INSIDE the engine container.

The method is the fleet's established one (the 2026-09-29 / 2026-10-01
tuning campaigns, /home/dank/mtp_measure.py): SSE chunk counting is invalid
under speculative decoding because rejected drafts change the accounting, so
decode is measured from the server's own usage counters:

    decode_tok_s = completion_tokens / (wall - ttft)

with ttft = wall time to the FIRST content-bearing chunk (thinking arrives in
`reasoning`/`delta.reasoning` on this stack and is excluded from the decode
window — and is deterministic per prompt at temperature 0, so the exclusion
cancels in any A/B on the same workload).

Workload (fixed, deterministic, offline):
  * prompts built from a large fixed-seed shuffled word list so tokens rarely
    repeat (no n-gram/prefix predictability of the text itself), sized with
    the server's own /v1/messages/count_tokens counter
  * depths 115 / 1056 / 4096 target prompt tokens
  * one 64-token warm pass per depth (warms the expert arena), then
    max_tokens=256 at temperature 0, 3 runs, median
Everything is greedy with a fixed seed — the same workload every run, no
network, no time-of-day dependence.

Env:
  AR_BASE             endpoint (default http://127.0.0.1:7770)
  AR_MODEL            served name (default frank-primary)
  AR_PROMPT_CEILING   hard prompt ceiling in tokens (default 100000)

Output: progress on stderr, ONE json object on stdout (the last line).
Exit codes: 0 ok, 2 endpoint unhealthy / no usable completion, 3 bad usage.
"""
import json
import os
import random
import statistics
import sys
import time
import urllib.error
import urllib.request

BASE = os.environ.get("AR_BASE", "http://127.0.0.1:7770")
MODEL = os.environ.get("AR_MODEL", "frank-primary")
PROMPT_CEILING = int(os.environ.get("AR_PROMPT_CEILING", "100000"))
DEPTHS = (115, 1056, 4096)
WARM_TOKENS = 64
MAX_TOKENS = 256
RUNS = 3

# 43-word x 12 = 516-word vocabulary, seed 1234 — the campaign's exact list
# and seed, so prompts are byte-identical to the historical measurements.
_VOCAB = ("alpha bravo charlie delta echo foxtrot golf hotel india juliet kilo lima "
          "mike november oscar papa quebec romeo sierra tango uniform victor whiskey "
          "xray yankee zulu cardinal ordinal topology kernel scheduler resident "
          "unified aperture expert arena slab offload prefetch gemm triton quant "
          "tensor shard checkpoint manifest embed router cache page table window "
          "gradient sparse dense moe attention head dim block stream decode prefill "
          "bandwidth latency throughput residency eviction tier swap pin memlock").split()
_R = random.Random(1234)
_ORDER = list(range(len(_VOCAB)))
_R.shuffle(_ORDER)


def log(msg):
    print("[measure] %s" % msg, file=sys.stderr, flush=True)


def post_json(path, body, timeout):
    req = urllib.request.Request(
        BASE + path, data=json.dumps(body).encode(),
        headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.load(r)


def count_tokens(text):
    """Token count without a prefill, from the server's own counter."""
    body = {"model": MODEL,
            "messages": [{"role": "user", "content": text}]}
    r = post_json("/v1/messages/count_tokens", body, timeout=300)
    n = r.get("input_tokens") or r.get("count") or r.get("token_count")
    if not n:
        raise RuntimeError("count_tokens returned no count: %r" % (r,))
    return int(n)


def words(n):
    return " ".join(_VOCAB[_ORDER[i % len(_VOCAB)]] for i in range(n))


def calibrate():
    """Measure words->tokens from the server's own counter."""
    return 2000.0 / count_tokens(words(2000))


def prompt_for(target_tokens, wpt):
    n = max(1, int(target_tokens * wpt))
    text = words(n)
    pt = count_tokens(text)
    while pt > PROMPT_CEILING:
        n = max(1, int(n * (PROMPT_CEILING / pt) * 0.98))
        text = words(n)
        pt = count_tokens(text)
    return text, pt


def measured(prompt, max_tokens):
    """Return (ttft_s, decode_tok_s, completion_tokens) via usage counters."""
    body = {"model": MODEL,
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": max_tokens, "temperature": 0,
            "stream": True, "stream_options": {"include_usage": True}}
    req = urllib.request.Request(
        BASE + "/v1/chat/completions", data=json.dumps(body).encode(),
        headers={"Content-Type": "application/json"})
    t0 = time.time()
    ttft = None
    usage = {}
    with urllib.request.urlopen(req, timeout=1800) as r:
        for raw in r:
            line = raw.decode("utf-8", "replace").strip()
            if not line.startswith("data:"):
                continue
            s = line[5:].strip()
            if s == "[DONE]":
                break
            try:
                ev = json.loads(s)
            except json.JSONDecodeError:
                continue
            if ev.get("usage"):
                usage = ev["usage"]
            for ch in ev.get("choices", []):
                if (ch.get("delta") or {}).get("content"):
                    if ttft is None:
                        ttft = time.time() - t0
    wall = time.time() - t0
    ct = int(usage.get("completion_tokens") or 0)
    if ct <= 0:
        # No usage on the stream: the endpoint is unhealthy or the model
        # returned nothing (e.g. HTTP 200 with an error payload).
        raise RuntimeError("no completion_tokens in usage: %r" % (usage,))
    dec_window = max(wall - (ttft or 0.0), 1e-9)
    return (ttft or wall), (ct / dec_window), ct


def main():
    log("base=%s model=%s ceiling=%d" % (BASE, MODEL, PROMPT_CEILING))
    try:
        wpt = calibrate()
    except Exception as e:
        log("endpoint unhealthy during calibration: %s" % e)
        return 2
    log("words/token: %.3f" % wpt)

    result = {"words_per_token": round(wpt, 5)}
    for tgt in DEPTHS:
        try:
            p, pt = prompt_for(tgt, wpt)
            measured(p, WARM_TOKENS)  # warm the arena
        except Exception as e:
            log("warm pass failed at depth %d: %s" % (tgt, e))
            return 2
        rates, tts, cts = [], [], []
        for i in range(RUNS):
            try:
                tt, r, ct = measured(p, MAX_TOKENS)
            except Exception as e:
                log("run %d/%d failed at depth %d: %s" % (i + 1, RUNS, tgt, e))
                continue
            rates.append(r)
            tts.append(tt)
            cts.append(ct)
        if len(rates) < 2:
            log("only %d usable run(s) at depth %d — refusing to median" % (len(rates), tgt))
            return 2
        result["decode_tps_%d" % tgt] = round(statistics.median(rates), 2)
        result["ttft_ms_%d" % tgt] = round(1000.0 * statistics.median(tts), 1)
        result["completion_tokens_%d" % tgt] = int(statistics.median(cts))
        result["prompt_tokens_%d" % tgt] = pt
        log("depth %5d tok  decode %.2f t/s  ttft %.3f s  (runs=%d)"
            % (pt, result["decode_tps_%d" % tgt], result["ttft_ms_%d" % tgt] / 1000.0, len(rates)))

    meds = [result["decode_tps_%d" % t] for t in DEPTHS]
    result["decode_tps"] = round(statistics.median(meds), 2)
    print(json.dumps(result))
    return 0


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