#!/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 VENV_PYTHON = sys.executable # use whichever python ran this script 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()