expert_profile_drive.py

#!/usr/bin/env python3
"""Drive coding traffic through a ds4-server running with DS4_EXPERT_PROFILE, then
(parse mode) summarize the routed-expert concentration / cache-hit curve it wrote.

Usage:
  expert_profile_drive.py drive  <server_url> <n_total>   # POST coding prompts, discard output
  expert_profile_drive.py parse  <profile.json>           # summarize concentration

The point: coder.dat (the imatrix) is COUNT-NORMALIZED so it measures intensity-when-routed,
not routing frequency, and reports near-uniform expert energy. The --expert-profile JSON
records hist (routing counts) + weight_hist (router gate mass) + an LRU cache simulation at
caps 1..384 — the right signal for 'how many experts does coding actually need resident'.
"""
import json, sys, urllib.request, time, os

DATA = os.path.expanduser("~/Beep/benchy-dash/data")
# representative coding mix (same prompt wrapper as eval_code.py), ~1h on Flash q2 streaming
MIX = [("bigcodebench", 40), ("bigcodebench_ext", 80), ("lcb_v6_func", 40), ("mbppplus", 20)]
WRAP = ("Complete the following Python task. Respond with ONLY the function "
        "implementation inside a single ```python code block — no prose, no examples.\n\n")


def load_prompts(n_total):
    out = []
    for bench, n in MIX:
        rows = [json.loads(l) for l in open(f"{DATA}/{bench}.jsonl") if l.strip()][:n]
        for r in rows:
            out.append((bench, WRAP + r["prompt"]))
    return out[:n_total] if n_total else out


def drive(url, n_total):
    prompts = load_prompts(n_total)
    print(f"driving {len(prompts)} coding prompts -> {url}", flush=True)
    t0 = time.time()
    for i, (bench, prompt) in enumerate(prompts, 1):
        body = json.dumps({"messages": [{"role": "user", "content": prompt}],
                           "max_tokens": 1536, "temperature": 0.0, "think": False}).encode()
        req = urllib.request.Request(url + "/v1/chat/completions", data=body,
                                     headers={"Content-Type": "application/json"})
        try:
            with urllib.request.urlopen(req, timeout=600) as resp:
                json.loads(resp.read())  # discard; we only want the routing it drove
        except Exception as e:
            print(f"  [{i}/{len(prompts)}] {bench} ERROR {e}", flush=True)
            continue
        if i % 10 == 0 or i == len(prompts):
            dt = time.time() - t0
            print(f"  [{i}/{len(prompts)}] {bench}  {dt:.0f}s  {dt/i:.1f}s/task", flush=True)
    print(f"drive done in {time.time()-t0:.0f}s", flush=True)


def parse(path):
    d = json.load(open(path))
    cs = {row["n"]: row for row in d.get("cache_summary", [])}
    layers = d.get("layers_detail", [])
    n_expert = d.get("experts") or 256
    uniq = sorted(l.get("unique_experts", 0) for l in layers)
    print(f"model={d.get('model')}  layers={d.get('layers')}  n_expert={n_expert}")
    print(f"total_selections={d.get('selections')}\n")
    print("LRU cache-hit curve (how much of the coding routing a K-expert/layer cache serves):")
    print(f"  {'K':>4}  {'hit_rate':>9}  {'weighted_hit_rate':>18}")
    for k in sorted(cs):
        r = cs[k]
        print(f"  {k:>4}  {r['hit_rate']*100:8.2f}%  {r['weighted_hit_rate']*100:17.2f}%")
    if uniq:
        import statistics
        print(f"\nunique experts routed / layer (of {n_expert}): "
              f"min {uniq[0]} median {int(statistics.median(uniq))} max {uniq[-1]}")
    # the headline numbers for the pruned-PRO question
    for k in (16, 24, 32):
        if k in cs:
            print(f"=> a {k}-expert/layer resident cache serves "
                  f"{cs[k]['weighted_hit_rate']*100:.1f}% of coding routing weight")


if __name__ == "__main__":
    if len(sys.argv) >= 2 and sys.argv[1] == "drive":
        drive(sys.argv[2].rstrip("/"), int(sys.argv[3]) if len(sys.argv) > 3 else 0)
    elif len(sys.argv) >= 3 and sys.argv[1] == "parse":
        parse(sys.argv[2])
    else:
        sys.exit(__doc__)