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, 3 runs, median — repeated for every temperature in AR_TEMPS
    (default 0 and 0.7). The prompt workload is deterministic (fixed seed, no
    network, no time-of-day dependence); at temperature > 0 the generated text
    is not, but the measured window is bounded by max_tokens, so the rate stays
    comparable and the non-determinism is exactly the sampling that a
    speculative-decode arm must survive.

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)
  AR_TEMPS            comma-separated sampling temperatures (default "0,0.7")

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
# Temperature sweep. The unsuffixed keys are temperature 0 — byte-comparable to
# every historical measurement — and AR_TEMPS adds more. The second one matters
# for speculative decode: draft ACCEPTANCE is sampling-dependent (a greedy draft
# matches a greedy target exactly and is rejected more under a non-greedy
# target), so an MTP arm judged only at temp 0 would be judged on half the
# contract. Production traffic is non-zero temperature (the payload's
# generation_config says 1.0; an override to 0.7 is under test).
TEMPS = tuple(float(x) for x in os.environ.get("AR_TEMPS", "0,0.7").split(","))

# 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, temperature=0):
    """Return (ttft_s, decode_tok_s, completion_tokens) via usage counters."""
    body = {"model": MODEL,
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": max_tokens, "temperature": temperature,
            "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 fmt_temp(t):
    return ("%.1f" % t).rstrip("0").rstrip(".") if t else "0"


def main():
    log("base=%s model=%s ceiling=%d temps=%s"
        % (BASE, MODEL, PROMPT_CEILING, ",".join(fmt_temp(t) for t in TEMPS)))
    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 and the prefix cache
        except Exception as e:
            log("warm pass failed at depth %d: %s" % (tgt, e))
            return 2
        for temp in TEMPS:
            suffix = "" if temp == 0 else "_t%s" % fmt_temp(temp)
            rates, tts, cts = [], [], []
            for i in range(RUNS):
                try:
                    tt, r, ct = measured(p, MAX_TOKENS, temp)
                except Exception as e:
                    log("run %d/%d failed at depth %d temp %s: %s"
                        % (i + 1, RUNS, tgt, fmt_temp(temp), e))
                    continue
                rates.append(r)
                tts.append(tt)
                cts.append(ct)
            if len(rates) < 2:
                log("only %d usable run(s) at depth %d temp %s — refusing to median"
                    % (len(rates), tgt, fmt_temp(temp)))
                return 2
            result["decode_tps_%d%s" % (tgt, suffix)] = round(statistics.median(rates), 2)
            result["ttft_ms_%d%s" % (tgt, suffix)] = round(1000.0 * statistics.median(tts), 1)
            result["completion_tokens_%d%s" % (tgt, suffix)] = int(statistics.median(cts))
            result["prompt_tokens_%d%s" % (tgt, suffix)] = pt
            log("depth %5d tok temp %-4s decode %7.2f t/s  ttft %.3f s  (runs=%d)"
                % (pt, fmt_temp(temp), result["decode_tps_%d%s" % (tgt, suffix)],
                   result["ttft_ms_%d%s" % (tgt, suffix)] / 1000.0, len(rates)))

    for temp in TEMPS:
        suffix = "" if temp == 0 else "_t%s" % fmt_temp(temp)
        meds = [result["decode_tps_%d%s" % (t, suffix)] for t in DEPTHS]
        result["decode_tps%s" % suffix] = round(statistics.median(meds), 2)
    print(json.dumps(result))
    return 0


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