harness/phase4_wallclock.py

"""Phase 4.1 — deployed wall-clock speedup via a faithful external spec-decode
loop with REAL timers (not engine-integrated).

Wiring a retrieval draft into vLLM/sglang's internal speculative-decoding hook
is out of reach in the available time, so — as the directive permits — we
reproduce the accept/reject/re-decode loop end-to-end against the live served
model with real wall-clock timers, and label it plainly as an external harness.

Spec-decode accounting (single retrieval draft per call): the target model runs
ONE verification forward over the drafted tool call, accepts its ``L``-token
correct prefix (token-LCP against the greedy target, exactly the MAT metric),
then autoregressively decodes the remaining ``T-L`` target tokens. So the target
performs ~``(T-L)`` sequential forwards plus one verify, vs. ``T`` for a
no-speculation baseline. We MEASURE the real per-request wall-clock of generating
``T`` tokens (baseline) and ``T-L`` tokens (each arm) from the actual decision-
point prompt on the served gpt-oss-120b — the content is irrelevant, only the
decode-step count and real server timing matter — and add one measured verify
forward. Reports p50/p95 latency per arm and end-to-end speedup vs no-memory.

Run from the repo root:  python -m harness.phase4_wallclock --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 .data import Task
from .memory import Embedder, NoMemory, PersonalMemory, StaticGlobal
from .run_accept import _parse_target
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")


def _collect_points(domains):
    """Replay the live tau2 decision points; record, per post-warmup point,
    each arm's token-accept length against the served target."""
    dp = [json.loads(l) for l in
          (RESULTS / "tau2_live_decision_points.jsonl").read_text().splitlines()
          if json.loads(l)["domain"] in domains]
    tools = {d: json.loads((ROOT / "data" / "tau2" /
                            f"tools_{d}.json").read_text()) for d in domains}
    tasks = [Task(id=r["id"], query=r["query"], functions=tools[r["domain"]],
                  origin_id=r["id"]) for r in dp]
    targets = {r["query"]: r["target"] for r in dp}
    emb = Embedder()
    inst = build_users(tasks, n_users=40, tasks_per_user=15, n_sessions=12,
                       queries_per_session=6, seed=0, perturb_prob=0.0)
    inst.sort(key=lambda x: (x.session, x.user_id))
    arms = [NoMemory(), StaticGlobal(), PersonalMemory(capacity=48, eviction="lru")]
    cur, pts = -1, []
    for ins in inst:
        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, "target": tgt,
                   "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)
    return pts


def _timed_gen(url, prompt, max_tok):
    """Real wall-clock (ms) to generate exactly max_tok tokens; content unused."""
    t0 = time.perf_counter()
    r = requests.post(url, json={"model": "gpt-oss-120b", "temperature": 0.0,
                                 "max_tokens": max_tok, "ignore_eos": True,
                                 "messages": [{"role": "user", "content": prompt}]},
                      timeout=180)
    dt = (time.perf_counter() - t0) * 1000
    r.raise_for_status()
    return dt


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--url", default="http://localhost:30000/v1")
    p.add_argument("--domains", nargs="+",
                   default=["airline", "retail", "telecom"])
    p.add_argument("--sample", type=int, default=120)
    p.add_argument("--seed", type=int, default=0)
    args = p.parse_args()
    chat = args.url.rstrip("/") + "/chat/completions"

    metrics.get_tokenizer(MODEL_PATH)
    pts = _collect_points(args.domains)
    rng = random.Random(args.seed)
    rng.shuffle(pts)
    pts = pts[: args.sample]
    print(f"[wallclock] {len(pts)} sampled decision points", flush=True)

    # one measured verify forward (prefill+1 tok) as spec-decode overhead
    verify_ms = sorted(_timed_gen(chat, pts[i]["query"][:1500], 1)
                       for i in range(min(15, len(pts))))
    verify = verify_ms[len(verify_ms) // 2]

    lat = defaultdict(list)      # arm -> per-request wall-clock ms
    for k, row in enumerate(pts):
        prompt = row["query"][:1500]
        T = max(1, row["T"])
        # cache decode-time per distinct token count to save calls
        need = {T}
        for arm in ("no_memory", "static_global", "personal_memory"):
            need.add(max(1, T - row[arm]))
        tcache = {n: _timed_gen(chat, prompt, n) for n in need}
        lat["baseline_no_spec"].append(tcache[T])
        for arm in ("no_memory", "static_global", "personal_memory"):
            lat[arm].append(verify + tcache[max(1, T - row[arm])])
        if (k + 1) % 20 == 0:
            print(f"  {k+1}/{len(pts)}", flush=True)

    def stats(xs):
        xs = sorted(xs)
        return {"p50_ms": round(xs[len(xs) // 2], 1),
                "p95_ms": round(xs[int(0.95 * len(xs))], 1),
                "mean_ms": round(sum(xs) / len(xs), 1)}

    base = stats(lat["baseline_no_spec"])
    out = {"config": {"domains": args.domains, "n": len(pts),
                      "verify_forward_ms": round(verify, 1),
                      "note": "faithful external spec-decode loop, real timers, "
                              "NOT engine-integrated (see docstring)"},
           "baseline_no_spec": base, "arms": {}}
    for arm in ("no_memory", "static_global", "personal_memory"):
        s = stats(lat[arm])
        s["speedup_vs_no_memory_p50"] = round(
            stats(lat["no_memory"])["p50_ms"] / s["p50_ms"], 3)
        s["speedup_vs_baseline_p50"] = round(base["p50_ms"] / s["p50_ms"], 3)
        out["arms"][arm] = s
    (RESULTS / "phase4_wallclock_deployed.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))


if __name__ == "__main__":
    main()