harness/run_ablation.py

"""Ablation: decompose PersonalMemory's gain over StaticGlobal into the
contributions of (a) per-user personalization and (b) online-growth+eviction.

2x2 design (personalized? x evicting/growing?):
  static_global    [-pers, -grow]  frozen global datastore (ToolSpec-style)  [baseline]
  global_evict     [-pers, +grow]  one global store, keeps ingesting + LRU-evicts
  personal_noevict [+pers, -evict] per-user store, grows online, UNBOUNDED
  personal_memory  [+pers, +evict] per-user store, grows online + LRU-evicts  [ours]

All four share the identical target stream (arm-independent greedy target decode
from the served model), so they differ ONLY in memory policy. Targets are cached
to disk so re-runs need no GPU. Single seed at full scale: the primary 3-seed run
established seed-std <= 0.20 MAT, so one seed is sufficient to attribute the gap.
"""
from __future__ import annotations

import argparse
import json
from collections import defaultdict
from pathlib import Path

from . import metrics
from .client import ToolClient
from .data import load_bfcl
from .memory import (Embedder, StaticGlobal, GlobalEvict, PersonalNoEvict,
                     PersonalMemory)
from .run_accept import generate_targets, _parse_target, MODEL_PATH
from .simulate import build_users

ROOT = Path(__file__).resolve().parent.parent
RESULTS = ROOT / "results"


def _make_arms(args):
    return [
        StaticGlobal(),
        GlobalEvict(capacity=args.users * args.capacity),
        PersonalNoEvict(),
        PersonalMemory(capacity=args.capacity, eviction="lru"),
    ]


def _replay(instances, targets, embedder, args):
    arms = _make_arms(args)
    agg = {a.name: defaultdict(list) for a in arms}
    cur_session = -1
    for ins in instances:
        tgt = targets.get(ins.query)
        if tgt is None:
            continue
        if ins.session != cur_session:
            cur_session = ins.session
            for a in arms:
                if isinstance(a, StaticGlobal) and cur_session == 1:
                    a.freeze()
        for a in arms:
            draft = a.draft(ins.query, ins.functions, ins.user_id, embedder)
            agg[a.name][ins.session].append(metrics.score(draft, tgt))
        name, argd = _parse_target(tgt)
        for a in arms:
            a.observe(ins.query, ins.functions, ins.user_id, name, argd, embedder)
            if isinstance(a, PersonalMemory) and ins.session == 0:
                a.seed_shared(ins.query, name, argd, embedder)
    return agg


def _overall(agg, warmup=0):
    out = {}
    for arm, per_sess in agg.items():
        scores = [x for s, xs in per_sess.items() if s > warmup for x in xs]
        n = len(scores)
        out[arm] = {
            "n": n,
            "MAT": round(sum(x["accept_length"] for x in scores) / n, 3),
            "accepted_frac": round(sum(x["accepted_frac"] for x in scores) / n, 4),
            "exact_rate": round(sum(1 for x in scores if x["exact"]) / n, 4),
        }
    return out


def _by_session(agg, n_sessions):
    out = {}
    for arm, per_sess in agg.items():
        out[arm] = {}
        for s in range(n_sessions):
            xs = per_sess.get(s, [])
            if xs:
                out[arm][str(s)] = round(
                    sum(x["accept_length"] for x in xs) / len(xs), 3)
    return out


def run(args):
    RESULTS.mkdir(exist_ok=True)
    metrics.get_tokenizer(args.model_path or MODEL_PATH)
    tasks = load_bfcl()
    embedder = Embedder()

    instances = build_users(
        tasks, n_users=args.users, tasks_per_user=args.tasks_per_user,
        n_sessions=args.sessions, queries_per_session=args.queries_per_session,
        seed=args.seed, perturb_prob=args.perturb_prob,
        arrival=args.arrival, novel_weight=args.novel_weight,
        warmup_frac=args.warmup_frac)
    instances.sort(key=lambda x: (x.session, x.user_id))

    cache = RESULTS / f"{args.tag}_targets_seed{args.seed}.json"
    if cache.exists():
        print(f"[ablation] loading cached targets from {cache.name}", flush=True)
        targets = json.loads(cache.read_text())
    else:
        client = ToolClient(url=args.url, model=args.model)
        if not client.ping():
            raise SystemExit(f"served model not reachable at {args.url}")
        print(f"[ablation] {len(instances)} instances; generating targets ...",
              flush=True)
        targets = generate_targets(client, instances, workers=args.workers)
        cache.write_text(json.dumps(targets, indent=1))
        print(f"[ablation] cached {len(targets)} targets -> {cache.name}",
              flush=True)

    n_none = sum(1 for v in targets.values() if v is None)
    agg = _replay(instances, targets, embedder, args)
    overall = _overall(agg, warmup=0)
    by_session = _by_session(agg, args.sessions)

    static = overall["static_global"]["MAT"]
    ours = overall["personal_memory"]["MAT"]
    gevict = overall["global_evict"]["MAT"]
    pnoev = overall["personal_noevict"]["MAT"]
    # Clean single-variable contrasts (each isolates ONE knob, holding the rest):
    #   live/growth : freeze->live on a global store (evict held on)   = gevict - static
    #   personalization: global->per-user, holding live+evict          = ours   - gevict
    #   eviction    : unbounded->bounded, holding live+per-user         = ours   - pnoev
    decomp = {
        "total_gap_ours_vs_static": round(ours - static, 3),
        "live_growth_effect (global_evict - static)": round(gevict - static, 3),
        "personalization_effect (personal_memory - global_evict)": round(ours - gevict, 3),
        "eviction_effect (personal_memory - personal_noevict)": round(ours - pnoev, 3),
        "live_growth_share_pct": round(100 * (gevict - static) / (ours - static), 1)
            if ours != static else None,
    }

    out = {
        "config": vars(args),
        "seed": args.seed,
        "n_instances": len(instances),
        "n_unique_queries": len(targets),
        "n_no_toolcall": n_none,
        "overall_post_warmup": overall,
        "by_session": by_session,
        "decomposition": decomp,
    }
    (RESULTS / f"{args.tag}_ablation_results.json").write_text(
        json.dumps(out, indent=2))

    print("\n=== Ablation: post-warmup MAT (sessions >= 1) ===", flush=True)
    for arm in ["static_global", "global_evict", "personal_noevict",
                "personal_memory"]:
        v = overall[arm]
        print(f"  {arm:>16}: MAT={v['MAT']:6.2f}  acc_frac={v['accepted_frac']:.3f}"
              f"  exact={v['exact_rate']:.3f}  n={v['n']}", flush=True)
    print("\n=== Decomposition ===", flush=True)
    for k, v in decomp.items():
        print(f"  {k}: {v}", flush=True)
    print(f"\nWrote results/{args.tag}_ablation_results.json", flush=True)


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--users", type=int, default=40)
    p.add_argument("--tasks-per-user", type=int, default=15)
    p.add_argument("--sessions", type=int, default=12)
    p.add_argument("--queries-per-session", type=int, default=6)
    p.add_argument("--capacity", type=int, default=48)
    p.add_argument("--workers", type=int, default=16)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--url", default="http://localhost:30000/v1")
    p.add_argument("--model", default="gpt-oss-120b")
    p.add_argument("--model-path", default="")
    p.add_argument("--tag", default="phase2")
    p.add_argument("--arrival", default="spread", choices=["spread","burst","late"])
    p.add_argument("--novel-weight", type=float, default=3.0)
    p.add_argument("--warmup-frac", type=float, default=0.5)
    p.add_argument("--perturb-prob", type=float, default=1.0,
                   help="probability a post-warmup query gets numeric-variant "
                        "perturbation (1.0 = legacy behaviour; 0.0 = verbatim "
                        "repeats only, isolating novel-task drift).")
    run(p.parse_args())


if __name__ == "__main__":
    main()