"""Phase 4 MAIN TABLE (Part A) — 4 arms x 4 real datasets, full metrics. Rows/arms (reuse existing, no reimplementation): no_memory = Vanilla AR (speedup reference) toolspec = faithful ToolSpec (Phase 4.4: confidence-gated retrieval + FSM) static_global = ToolSpec + static memory (simple frozen proxy) personal_memory= SpecMem (ours: live, evicting, per-user) Columns/datasets (all REAL): API-Bank, ToolAlpaca, BFCLv4, ToolBench. Per (arm, dataset) cell we record: 1. MAT (mean accepted tokens) -- replay 2. tokens/s (real decode throughput) -- wall-clock timing 3. e2e wall-clock speedup vs Vanilla AR -- wall-clock timing (p50) 4. retrieval / write-back overhead ms -- timed embed+lookup / embed+insert 5. memory size (entries) at run end -- replay 6. cold-start vs long-term MAT curve -- per-session MAT (replay) Targets are the served gpt-oss-120b's greedy calls (real, cached per dataset+seed). Run from the repo root: python -m harness.phase4_maintable --url http://localhost:30000/v1 """ from __future__ import annotations import argparse import json import os import random import time from collections import defaultdict from pathlib import Path import requests from . import metrics from .client import ToolClient from .data import (load_apibank, load_bfcl, load_toolalpaca, load_toolbench) from .memory import (Embedder, NoMemory, PersonalMemory, StaticGlobal, ToolSpecBaseline) from .run_accept import _parse_target, generate_targets from .simulate import build_users ROOT = Path(__file__).resolve().parent.parent RESULTS = ROOT / "results" # Tokenizer for the token-LCP accept metric: HF hub id by default; # override with a local snapshot path if running offline. MODEL_PATH = os.environ.get("SPECMEM_TOKENIZER", "openai/gpt-oss-120b") DATASETS = {"apibank": load_apibank, "toolalpaca": load_toolalpaca, "bfcl": load_bfcl, "toolbench": load_toolbench} ARMS = ["no_memory", "toolspec", "static_global", "personal_memory"] def _make_arms(cap): return [NoMemory(), ToolSpecBaseline(), StaticGlobal(), PersonalMemory(capacity=cap, eviction="lru")] def _mem_size(arm): if isinstance(arm, PersonalMemory): return sum(len(s) for s in arm.stores.values()) if isinstance(arm, (StaticGlobal, ToolSpecBaseline)): return len(arm.entries) return 0 def _replay(instances, targets, emb, cap): """3-arm+baseline replay; return per-arm MAT, per-session MAT, mem size.""" arms = _make_arms(cap) agg = {a.name: defaultdict(list) for a in arms} cur = -1 for ins in instances: tgt = targets.get(ins.query) if tgt is None: continue if ins.session != cur: cur = ins.session if cur == 1: for a in arms: if hasattr(a, "freeze"): a.freeze() for a in arms: agg[a.name][ins.session].append(metrics.score( a.draft(ins.query, ins.functions, ins.user_id, emb), tgt)) cn, ca = _parse_target(tgt) for a in arms[1:]: a.observe(ins.query, ins.functions, ins.user_id, cn, ca, emb) if isinstance(a, PersonalMemory) and ins.session == 0: a.seed_shared(ins.query, cn, ca, emb) out = {} for a in arms: v = agg[a.name] post = [x for s, xs in v.items() if s > 0 for x in xs] by_sess = {s: round(sum(x["accept_length"] for x in xs) / len(xs), 3) for s, xs in sorted(v.items()) if xs} out[a.name] = { "MAT": round(sum(x["accept_length"] for x in post) / max(1, len(post)), 3), "accepted_frac": round(sum(x["accepted_frac"] for x in post) / max(1, len(post)), 3), "by_session": by_sess, "mem_entries": _mem_size(a)} return out def _timed_gen(chat, prompt, ntok): t0 = time.perf_counter() r = requests.post(chat, json={"model": "gpt-oss-120b", "temperature": 0.0, "max_tokens": ntok, "ignore_eos": True, "messages": [{"role": "user", "content": prompt}]}, timeout=180) dt = (time.perf_counter() - t0) * 1000 r.raise_for_status() return dt def _wallclock(instances, targets, emb, cap, chat, sample, seed): """Real spec-decode wall-clock per arm on a sample: p50/p95 + speedup + tokens/s. Rebuild arm accept-lengths at each sampled point via replay.""" arms = _make_arms(cap) pts, cur = [], -1 for ins in instances: tgt = targets.get(ins.query) if tgt is None: continue if ins.session != cur: cur = ins.session if cur == 1: for a in arms: if hasattr(a, "freeze"): a.freeze() if ins.session > 0: row = {"query": ins.query, "T": metrics.accept_length(tgt, tgt)[1]} for a in arms: row[a.name] = metrics.accept_length( a.draft(ins.query, ins.functions, ins.user_id, emb), tgt)[0] pts.append(row) cn, ca = _parse_target(tgt) for a in arms[1:]: a.observe(ins.query, ins.functions, ins.user_id, cn, ca, emb) if isinstance(a, PersonalMemory) and ins.session == 0: a.seed_shared(ins.query, cn, ca, emb) rng = random.Random(seed) rng.shuffle(pts) pts = [p for p in pts if p["T"] >= 2][:sample] verify = sorted(_timed_gen(chat, pts[i]["query"][:1500], 1) for i in range(min(12, len(pts))))[6 // 2 or 0] lat = defaultdict(list) toks = defaultdict(list) for row in pts: prompt, T = row["query"][:1500], max(1, row["T"]) need = {T} | {max(1, T - row[a]) for a in ARMS} tc = {n: _timed_gen(chat, prompt, n) for n in need} lat["baseline"].append(tc[T]) for a in ARMS: ms = verify + tc[max(1, T - row[a])] lat[a].append(ms) toks[a].append(T / (ms / 1000.0)) # effective tokens/s def p(xs, q): xs = sorted(xs) return xs[min(len(xs) - 1, int(q * len(xs)))] base_p50 = p(lat["baseline"], 0.5) out = {} for a in ARMS: out[a] = {"p50_ms": round(p(lat[a], 0.5), 1), "p95_ms": round(p(lat[a], 0.95), 1), "tokens_per_s": round(sum(toks[a]) / len(toks[a]), 1), "speedup_vs_vanilla": round( p(lat["no_memory"], 0.5) / p(lat[a], 0.5), 3)} out["_baseline_p50_ms"] = round(base_p50, 1) return out def _overhead(emb, chat): """Real retrieval (embed+NN over 48-entry store) and write-back (embed+ insert) latency, ms/query, measured separately.""" from .memory import Entry, _best_match import numpy as np store = [Entry(emb.embed(f"seed query {i}"), "x") for i in range(48)] rlat, wlat = [], [] for i in range(150): q = f"overhead probe query number {i} with args {i*7}" t0 = time.perf_counter(); e = emb.embed(q); _best_match(e, store) rlat.append((time.perf_counter() - t0) * 1000) t0 = time.perf_counter(); e2 = emb.embed(q + " wb"); store.append(Entry(e2, "y")) wlat.append((time.perf_counter() - t0) * 1000) return {"retrieval_ms_mean": round(sum(rlat) / len(rlat), 2), "writeback_ms_mean": round(sum(wlat) / len(wlat), 2)} def main(): p = argparse.ArgumentParser() p.add_argument("--url", default="http://localhost:30000/v1") p.add_argument("--model", default="gpt-oss-120b") p.add_argument("--datasets", nargs="+", default=list(DATASETS)) p.add_argument("--seeds", nargs="+", type=int, default=[0, 1, 2]) p.add_argument("--wallclock-sample", type=int, default=80) p.add_argument("--tasks-per-user", type=int, default=10) p.add_argument("--capacity", type=int, default=48) p.add_argument("--workers", type=int, default=16) args = p.parse_args() chat = args.url.rstrip("/") + "/chat/completions" metrics.get_tokenizer(MODEL_PATH) emb = Embedder() client = ToolClient(url=args.url, model=args.model) assert client.ping(), f"served model not reachable at {args.url}" overhead = _overhead(emb, chat) table = {} for ds in args.datasets: tasks = DATASETS[ds]() n_users = min(40, len(tasks) // args.tasks_per_user) cfg = {"n_users": n_users, "tasks_per_user": args.tasks_per_user, "pool": len(tasks)} print(f"\n=== {ds}: {len(tasks)} tasks -> {n_users} users x " f"{args.tasks_per_user} ===", flush=True) per_seed = defaultdict(dict) by_session_acc = defaultdict(lambda: defaultdict(list)) for sd in args.seeds: inst = build_users(tasks, n_users=n_users, tasks_per_user=args.tasks_per_user, n_sessions=12, queries_per_session=6, seed=sd) inst.sort(key=lambda x: (x.session, x.user_id)) cache_f = RESULTS / f"phase4_mt_targets_{ds}_seed{sd}.json" cache = json.loads(cache_f.read_text()) if cache_f.exists() else {} fmap = {i.query: i.functions for i in inst} miss = [type("S", (), {"query": q, "functions": fmap[q]})() for q in {i.query for i in inst} if q not in cache] if miss: print(f" [seed {sd}] generating {len(miss)} targets ...", flush=True) cache.update(generate_targets(client, miss, workers=args.workers)) cache_f.write_text(json.dumps(cache)) res = _replay(inst, cache, emb, args.capacity) for a in ARMS: per_seed[a][sd] = res[a]["MAT"] for s, m in res[a]["by_session"].items(): by_session_acc[a][s].append(m) if sd == args.seeds[0]: mem = {a: res[a]["mem_entries"] for a in ARMS} wc = _wallclock(inst, cache, emb, args.capacity, chat, args.wallclock_sample, sd) print(f" [seed {sd}] MAT: " + " ".join(f"{a}={res[a]['MAT']}" for a in ARMS), flush=True) import statistics as st cell = {} for a in ARMS: mats = [per_seed[a][sd] for sd in args.seeds] cell[a] = { "MAT_mean": round(st.mean(mats), 3), "MAT_std": round(st.pstdev(mats), 3), "tokens_per_s": wc[a]["tokens_per_s"], "speedup_vs_vanilla": wc[a]["speedup_vs_vanilla"], "wallclock_p50_ms": wc[a]["p50_ms"], "wallclock_p95_ms": wc[a]["p95_ms"], "mem_entries": mem[a], "by_session_MAT": {s: round(sum(v) / len(v), 3) for s, v in sorted(by_session_acc[a].items())}, } table[ds] = {"config": cfg, "baseline_p50_ms": wc["_baseline_p50_ms"], "cells": cell} out = {"overhead": overhead, "arms": ARMS, "datasets": list(args.datasets), "table": table, "note": ("wall-clock via faithful external spec-decode loop, real " "sglang timers (ignore_eos), NOT engine-integrated; targets " "= served gpt-oss-120b greedy; BFCL is v4 (superset of the " "v2 ToolSpec used); ToolBench from OpenBMB Drive.")} (RESULTS / "phase4_main_table.json").write_text(json.dumps(out, indent=2)) print("\n=== SPEEDUP vs Vanilla AR (p50) ===") hdr = "arm".ljust(16) + "".join(d[:9].ljust(11) for d in args.datasets) print(hdr) for a in ARMS: row = a.ljust(16) + "".join( f"{table[d]['cells'][a]['speedup_vs_vanilla']}x".ljust(11) for d in args.datasets) print(row) if __name__ == "__main__": main()