whetstone_DSL/specialists/scripts/capacity_sweep.py

#!/usr/bin/env python3
"""
capacity_sweep.py

Runs a staged capacity sweep for one specialist gate.
Tiers: small+ (256/2) → medium (384/3) → large (512/4).

For each tier:
  1. Run a short pilot (3000 steps). If accuracy doesn't improve over
     the previous tier's result, skip the full run.
  2. Run the full training (up to --max-steps) only if pilot improved.
  3. Evaluate with eval_specialist.py (confusion, confidence, thresholds).
  4. Append result row to the sweep table.

Usage:
    python3 specialists/scripts/capacity_sweep.py \\
        --name        worker_type \\
        --train       specialists/data/generated/worker_type_train.tsv \\
        --eval        specialists/data/generated/worker_type_eval.tsv \\
        --labels      "implementer,reviewer,architect,qa" \\
        --runs-dir    /mnt/storage/fabricate_runs \\
        --results-out specialists/eval/results/worker_type_sweep.json

    # Run from whetstone_DSL root.
"""

import argparse
import json
import os
import shutil
import subprocess
import sys
import time
from pathlib import Path

ROOT        = Path(__file__).parent.parent.parent   # whetstone_DSL root
FABRICATE   = Path("/home/bill/Documents/WhetstoneAI_Fabricate")
VENV_PYTHON   = FABRICATE / ".venv" / "bin" / "python3"
TRAIN_SCRIPT  = ROOT / "specialists" / "scripts" / "train_specialist_pt.py"
EVAL_SCRIPT   = ROOT / "specialists" / "eval" / "eval_specialist_pt.py"

TIERS = [
    {"name": "small_plus", "hidden": 256, "layers": 2, "heads": 4},
    {"name": "medium",     "hidden": 384, "layers": 3, "heads": 6},
    {"name": "large",      "hidden": 512, "layers": 4, "heads": 8},
]

PILOT_STEPS   = 3000   # steps for pilot run
FULL_STEPS    = 15000  # max steps for full run
BLOCK         = 1000   # steps per eval block
BATCH_SIZE    = 256
LR            = 0.001
WD            = 0.01


def log(msg):
    print(msg, flush=True)


def run_training(run_dir: Path, train_data: Path, eval_data: Path,
                 labels: str, hidden: int, layers: int, heads: int,
                 max_steps: int, reset: bool, history_out: Path) -> bool:
    """Invoke train_specialist_pt.py. Returns True on success."""
    run_dir.mkdir(parents=True, exist_ok=True)
    if reset:
        for f in run_dir.glob("*"):
            f.unlink()
        history_out.unlink(missing_ok=True)

    cmd = [
        str(VENV_PYTHON), str(TRAIN_SCRIPT),
        "--dataset",               str(train_data),
        "--eval-dataset",          str(eval_data),
        "--out-dir",               str(run_dir),
        "--labels",                labels,
        "--lr",                    str(LR),
        "--weight-decay",          str(WD),
        "--max-steps",             str(max_steps),
        "--hidden-dim",            str(hidden),
        "--layers",                str(layers),
        "--heads",                 str(heads),
        "--batch-size",            str(BATCH_SIZE),
        "--grok-loss-threshold",   "0.05",
        "--grok-acc-jump",         "10.0",
        "--stop-after-grokking-blocks", "4",
        "--history-out",           str(history_out),
        "--reset",
    ]

    result = subprocess.run(cmd, capture_output=False)
    return result.returncode == 0


def run_eval(run_dir: Path, eval_data: Path, labels: str) -> dict | None:
    """Run eval_specialist_pt.py on the checkpoint. Returns parsed JSON or None."""
    ck = run_dir / "checkpoint.pt"
    if not ck.exists():
        return None

    result = subprocess.run(
        [str(VENV_PYTHON), str(EVAL_SCRIPT),
         "--checkpoint", str(ck),
         "--dataset",    str(eval_data),
         "--labels",     labels,
         "--quiet"],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        log(f"  [WARN] eval failed: {result.stderr[-500:]}")
        return None
    try:
        # Last line should be the JSON dump
        last_line = result.stdout.strip().splitlines()[-1]
        return json.loads(last_line)
    except (json.JSONDecodeError, IndexError):
        log(f"  [WARN] could not parse eval JSON")
        return None


def read_history(history_path: Path) -> dict:
    if not history_path.exists():
        return {}
    return json.loads(history_path.read_text())


def get_best_accuracy(history: dict) -> float:
    entries = history.get("history", [])
    if not entries:
        return 0.0
    return max(e.get("overall", 0.0) for e in entries)


def suggest_action(eval_result: dict, prev_accuracy: float) -> str:
    if eval_result is None:
        return "training_failed"
    acc = eval_result.get("overall_accuracy", 0.0)
    verdict = eval_result.get("verdict", "unknown")

    delta = acc - prev_accuracy
    if delta < 1.0:
        return "no_gain_stop_scaling"
    return verdict


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--name",        required=True, help="Gate name, e.g. worker_type")
    ap.add_argument("--train",       required=True, help="Training TSV path")
    ap.add_argument("--eval",        required=True, help="Eval TSV path")
    ap.add_argument("--labels",      required=True, help="Comma-sep label tokens")
    ap.add_argument("--runs-dir",    default="/mnt/storage/fabricate_runs")
    ap.add_argument("--results-out", default=None,  help="JSON output path")
    ap.add_argument("--tiers",       default="small_plus,medium,large",
                    help="Comma-sep tier names to run (subset of small_plus,medium,large)")
    ap.add_argument("--pilot-only",  action="store_true",
                    help="Only run pilots, no full training runs")
    args = ap.parse_args()

    train_data  = Path(args.train)
    eval_data   = Path(args.eval)
    runs_base   = Path(args.runs_dir)
    active_tiers = [t for t in TIERS if t["name"] in args.tiers.split(",")]

    sweep_results = {
        "gate": args.name,
        "labels": args.labels,
        "train": str(train_data),
        "eval": str(eval_data),
        "tiers": [],
    }

    prev_accuracy = 0.0  # baseline for gain comparison
    results_dir = ROOT / "specialists" / "eval" / "results"
    results_dir.mkdir(parents=True, exist_ok=True)

    for tier in active_tiers:
        tname   = tier["name"]
        hidden  = tier["hidden"]
        layers  = tier["layers"]
        heads   = tier.get("heads", 4)
        run_id  = f"whetstone_{args.name}_{tname}"
        run_dir = runs_base / run_id
        history = run_dir / "history.json"
        tier_result_path = results_dir / f"{args.name}_{tname}.json"

        log(f"\n{'='*60}")
        log(f"  Tier: {tname}  hidden={hidden} layers={layers} heads={heads}")
        log(f"  Run dir: {run_dir}")
        log(f"{'='*60}")

        tier_row = {
            "tier": tname,
            "hidden": hidden,
            "layers": layers,
            "heads": heads,
            "model_size_approx": f"~{hidden*hidden*layers*12 // 1_000_000}M params",
        }

        # --- Pilot run ---
        log(f"\n  [pilot] {PILOT_STEPS} steps...")
        t0 = time.perf_counter()
        ok = run_training(run_dir, train_data, eval_data, args.labels,
                          hidden, layers, heads, PILOT_STEPS, reset=True, history_out=history)
        pilot_time = time.perf_counter() - t0
        pilot_history = read_history(history)
        pilot_acc = get_best_accuracy(pilot_history)

        log(f"  [pilot] done in {pilot_time:.0f}s  best={pilot_acc:.1f}%")

        if not ok or pilot_acc == 0.0:
            log(f"  [pilot] FAILED or no accuracy recorded — skipping tier")
            tier_row["pilot_accuracy"] = None
            tier_row["action"] = "pilot_failed"
            sweep_results["tiers"].append(tier_row)
            continue

        tier_row["pilot_accuracy"] = round(pilot_acc, 1)
        tier_row["pilot_time_s"] = round(pilot_time, 0)

        gain = pilot_acc - prev_accuracy
        log(f"  [pilot] gain vs prev tier: {gain:+.1f}pp")

        if gain < 1.0 and prev_accuracy > 0:
            log(f"  [pilot] <1pp gain — skipping full run for this tier")
            tier_row["action"] = "no_gain_stop_scaling"
            sweep_results["tiers"].append(tier_row)
            break

        if args.pilot_only:
            log(f"  [pilot-only mode] evaluating pilot checkpoint...")
        else:
            # --- Full run ---
            log(f"\n  [full] continuing to {FULL_STEPS} steps...")
            t1 = time.perf_counter()
            run_training(run_dir, train_data, eval_data, args.labels,
                         hidden, layers, heads, FULL_STEPS, reset=False, history_out=history)
            full_time = time.perf_counter() - t1
            full_history = read_history(history)
            full_acc = get_best_accuracy(full_history)
            log(f"  [full] done in {full_time:.0f}s  best={full_acc:.1f}%")
            tier_row["full_accuracy"] = round(full_acc, 1)
            tier_row["full_time_s"] = round(full_time, 0)

        # --- Eval ---
        log(f"\n  [eval] running extended evaluation...")
        eval_result = run_eval(run_dir, eval_data, args.labels)
        if eval_result:
            # Save per-tier eval result
            tier_result_path.write_text(json.dumps(eval_result, indent=2))
            log(f"  [eval] accuracy={eval_result['overall_accuracy']}%  "
                f"verdict={eval_result['verdict']}")

            tier_row["eval_accuracy"] = eval_result["overall_accuracy"]
            tier_row["per_class"] = eval_result.get("per_class", {})
            tier_row["confusion_summary"] = eval_result.get("confusion_summary", [])[:5]
            tier_row["threshold_analysis"] = eval_result.get("threshold_analysis", [])
            tier_row["guardrail_threshold"] = eval_result.get("guardrail_threshold")
            tier_row["guardrail_abstain_rate"] = eval_result.get("guardrail_abstain_rate")
            tier_row["eval_latency_ms"] = eval_result.get("eval_latency_ms_per_example")
            tier_row["verdict"] = eval_result["verdict"]
            tier_row["action"] = suggest_action(eval_result, prev_accuracy)
            tier_row["eval_result_path"] = str(tier_result_path)

            prev_accuracy = eval_result["overall_accuracy"]
        else:
            tier_row["action"] = "eval_failed"

        sweep_results["tiers"].append(tier_row)

    # --- Final table ---
    log(f"\n{'='*60}")
    log(f"  SWEEP RESULTS: {args.name}")
    log(f"{'='*60}")
    header = f"  {'tier':12s}  {'hidden/L':8s}  {'pilot':7s}  {'full':7s}  {'verdict'}"
    log(header)
    for row in sweep_results["tiers"]:
        pilot = f"{row.get('pilot_accuracy','—'):>5}" if row.get('pilot_accuracy') else "  —  "
        full  = f"{row.get('eval_accuracy','—'):>5}" if row.get('eval_accuracy') else "  —  "
        hl    = f"{row['hidden']}/{row['layers']}"
        log(f"  {row['tier']:12s}  {hl:8s}  {pilot}%  {full}%  {row.get('verdict','—')}")

    out = args.results_out or str(results_dir / f"{args.name}_sweep.json")
    Path(out).write_text(json.dumps(sweep_results, indent=2))
    log(f"\nFull results → {out}")


if __name__ == "__main__":
    main()