1"""Scheduler: assignment manifests (paired tasks × arms × repeats, randomised order, blocked by task and execution window),2and the per-run pipeline (prepare → launch → controller-enforced gate cycles (C/D) → separate review (D) → collect →3contract check → evaluate). Runs are executed concurrently across executor nodes with a per-node concurrency limit."""4from __future__ import annotations5import json, os, posixpath, random, shlex, threading, time, traceback, concurrent.futures as cf6from .contracts import RunManifest, Budget, ARMS, utcnow, sha256_file, sha256_text, new_run_id, new_session_id7from .budgets import Envelope8from . import launcher, isolation, provenance9from .evaluation import evaluate_run1011class Controller:12 def __init__(self, repo_root: str, store, config: dict, nodes: list[str], remote_root="/Users/simon-pierreboucher/wp29", log=print):13 self.repo = repo_root; self.store = store; self.cfg = config; self.nodes = nodes; self.remote_root = remote_root; self.log = log14 self.arms_dir = os.path.join(repo_root, "protocol", "arms")15 self._node_sems = {n: threading.Semaphore(int(config.get("concurrency_per_node", 2))) for n in nodes}16 self._rr = 0; self._rr_lock = threading.Lock()17 self.model = config.get("model", "claude-opus-5-5")18 self.effort = config.get("effort")19 self.budget = Budget(**config.get("budget", {}))20 self._halt = threading.Event()2122 # ---------- assignment23 def make_assignment(self, split: str, task_ids: list[str], repeats: int, seed: int, arms=ARMS, start_seq=1) -> list[RunManifest]:24 """Paired design: every task gets every arm × repeats. Order: tasks shuffled; within a task block, (arm, repeat) pairs shuffled.25 Blocks are interleaved so that no arm runs systematically earlier than another (temporal blocking)."""26 rng = random.Random(seed)27 tasks = list(task_ids); rng.shuffle(tasks)28 runs = []; seq = start_seq29 for bi, t in enumerate(tasks):30 cells = [(a, r) for a in arms for r in range(repeats)]; rng.shuffle(cells)31 for oi, (a, r) in enumerate(cells):32 runs.append(RunManifest(run_id=new_run_id(split, seq), task_id=t, split=split, treatment=a, repeat_index=r,33 block_id=f"{split}-B{bi+1:02d}", order_in_block=oi, model_requested=self.model,34 budget=self.budget.__dict__.copy(), notes=f"assignment seed {seed}"))35 seq += 136 return runs3738 def register(self, runs: list[RunManifest]) -> None:39 for m in runs:40 m.code_commit = provenance.code_commit(self.repo)41 self.store.create_run(m)4243 # ---------- helpers44 def _pick_node(self) -> str:45 with self._rr_lock:46 n = self.nodes[self._rr % len(self.nodes)]; self._rr += 1; return n4748 def _read(self, *parts) -> str:49 return open(os.path.join(self.repo, *parts), encoding="utf-8").read()5051 def _bundle_dir(self, task_id) -> str:52 return os.path.join(self.repo, "tasks", "public_bundles", task_id)5354 def _run_dir(self, m) -> str:55 return os.path.join(self.repo, "runs", m["split"] if isinstance(m, dict) else m.split, m["run_id"] if isinstance(m, dict) else m.run_id)5657 # ---------- pipeline58 def execute(self, run_id: str, fault=None) -> dict:59 """Run the whole pipeline for one registered run; always returns, never raises (errors are recorded)."""60 m = self.store.get_run(run_id)61 node = m.get("node") or self._pick_node()62 with self._node_sems[node]:63 try:64 return self._execute(run_id, node, fault)65 except Exception as e: # noqa: BLE00166 tb = traceback.format_exc()[-1500:]67 self.log(f"[{run_id}] INFRA ERROR: {e}")68 try:69 cur = self.store.get_run(run_id)["status"]70 if cur in ("not_started", "prepared", "running", "interrupted"):71 self.store.transition(run_id, "error", note=str(e)[:300], terminal_reason="infrastructure_error:" + str(e)[:200], ended_at_utc=utcnow())72 except Exception: pass # noqa: BLE00173 return {"run_id": run_id, "status": "error", "error": str(e), "traceback": tb}7475 def _execute(self, run_id: str, node: str, fault) -> dict:76 m = self.store.get_run(run_id); arm = m["treatment"]; task_id = m["task_id"]77 rd = self._run_dir(m); os.makedirs(rd, exist_ok=True)78 attempt = self.store.n_error_attempts(run_id)79 if not self.store.log_job(f"{run_id}:execute:attempt{attempt}", run_id, "execute", node):80 # duplicate submission: idempotency guard81 self.log(f"[{run_id}] duplicate execute request ignored (job already logged)")82 return {"run_id": run_id, "status": "duplicate_ignored"}83 # 1. prepare workspace84 bundle = self._bundle_dir(task_id)85 prep = isolation.prepare_workspace(node, self.remote_root, run_id, bundle)86 if not prep["verified_identical"]:87 raise RuntimeError(f"workspace manifest mismatch: extra={prep['remote_extra']} missing={prep['remote_missing']}")88 ws, ctl = prep["workspace"], prep["ctl_dir"]89 task_md = self._read("tasks", "public_bundles", task_id, "TASK.md")90 common = self._read("protocol", "arms", "common_preamble.md"); arm_text = self._read("protocol", "arms", f"arm_{arm}.md")91 task_prompt = self._read("protocol", "arms", "task_prompt.md")92 sys_append, ctx_hash = isolation.effective_context(common, arm_text, task_md, task_prompt)93 isolation.write_ctl_files(node, ctl, {"system_append.md": sys_append, "prompt.txt": task_prompt,94 "reviewer_system.md": self._read("protocol", "arms", "reviewer.md"),95 "reviewer_prompt.txt": "Review the deliverables in this workspace now, as instructed, and write out/review/review.json.\n"})96 env = provenance.node_environment(node, launcher.ssh)97 self.store.update_manifest(run_id, node=node, task_hash=sha256_file(os.path.join(bundle, "bundle.json")),98 protocol_hash=sha256_text(sys_append), effective_context_hash=ctx_hash,99 environment_id=env["environment_id"], client_version=env.get("claude"), model_requested=self.model,100 code_commit=provenance.code_commit(self.repo))101 json.dump({"workspace_manifest": prep["manifest"], "node_environment": env, "system_append_sha256": sha256_text(sys_append)}, open(os.path.join(rd, "prepare.json"), "w"), indent=1)102 open(os.path.join(rd, "system_append.md"), "w").write(sys_append)103 self.store.transition(run_id, "prepared", note=f"workspace {ws} on {node}")104 if fault and fault.get("type") == "missing_fixture":105 from . import recovery106 fr = recovery.remove_fixture(node, ws, fault["path"])107 self.store.add_fault(run_id, "missing_fixture", fault["path"], fr)108 self.store.update_manifest(run_id, fault={"fault_type": "missing_fixture", "path": fault["path"], "utc": fr["utc"], "removed": fr["ok"]})109 # 2. executor session(s)110 envl = Envelope(self.budget.max_budget_usd, self.budget.wall_seconds, self.budget.max_turns_per_session, min_continuation_usd=self.budget.min_continuation_usd)111 raw = os.path.join(rd, "raw_stream.jsonl")112 sid = new_session_id(); seq = 0; sessions = []; status = "completed"; reason = "executor_finished"113 self.store.transition(run_id, "running", note="first session", started_at_utc=utcnow(), session_ids=[sid])114 prompt_file = posixpath.join(ctl, "prompt.txt"); sys_file = posixpath.join(ctl, "system_append.md")115 fault_rec = None116 r, seq = self._session(run_id, node, ws, sid, prompt_file, sys_file, envl, raw, seq, role="executor", kind="initial", fault=fault)117 sessions.append(sid)118 if r.get("fault_fired"):119 fault_rec = r["fault_fired"]120 # standardized interruption → resume once from the same session121 self.store.transition(run_id, "interrupted", note="fault: " + fault_rec["fault_type"])122 if envl.can_continue():123 isolation.write_ctl_files(node, ctl, {"resume_prompt.txt": self._read("protocol", "arms", "resume_after_interruption_prompt.md")})124 self.store.transition(run_id, "running", note="resume after interruption")125 r, seq = self._session(run_id, node, ws, sid, posixpath.join(ctl, "resume_prompt.txt"), sys_file, envl, raw, seq, role="executor", kind="resume_after_fault", resume=True)126 status, reason = self._status_from(r, envl)127 # 3. controller-enforced gate cycles for C/D128 gate_receipts = []129 if arm in ("C", "D") and status == "completed":130 for cycle in range(1, self.budget.max_gate_cycles + 1):131 rec = self.remote_gate(node, ws)132 self.store.add_gate_receipt(run_id, cycle, "controller", rec); gate_receipts.append(rec)133 if rec.get("passed"):134 break135 if not envl.can_continue():136 reason = "gate_failed_budget_exhausted"; break137 fb = self._read("protocol", "arms", "gate_feedback_prompt.md") + self._gate_text(rec)138 isolation.write_ctl_files(node, ctl, {f"gate_feedback_{cycle}.txt": fb})139 r, seq = self._session(run_id, node, ws, sid, posixpath.join(ctl, f"gate_feedback_{cycle}.txt"), sys_file, envl, raw, seq, role="executor", kind=f"gate_cycle_{cycle}", resume=True)140 status, reason = self._status_from(r, envl)141 if status != "completed":142 break143 else:144 rec = self.remote_gate(node, ws); self.store.add_gate_receipt(run_id, self.budget.max_gate_cycles + 1, "controller_final", rec); gate_receipts.append(rec)145 # 4. separate review for D146 review = None147 if arm == "D" and status == "completed" and gate_receipts and gate_receipts[-1].get("passed") and envl.can_continue():148 rsid = new_session_id(); sessions.append(rsid)149 launcher.ssh(node, f"mkdir -p {shlex.quote(ws)}/out/review")150 rr, seq = self._session(run_id, node, ws, rsid, posixpath.join(ctl, "reviewer_prompt.txt"), posixpath.join(ctl, "reviewer_system.md"), envl, raw, seq, role="reviewer", kind="review")151 review = self._fetch_json(node, posixpath.join(ws, "out", "review", "review.json"))152 self.store.add_review(run_id, rsid, review)153 if review and review.get("verdict") == "revise" and envl.can_continue():154 fb = self._read("protocol", "arms", "review_feedback_prompt.md") + json.dumps(review, indent=1)[:6000]155 isolation.write_ctl_files(node, ctl, {"review_feedback.txt": fb})156 r, seq = self._session(run_id, node, ws, sid, posixpath.join(ctl, "review_feedback.txt"), sys_file, envl, raw, seq, role="executor", kind="review_repair", resume=True)157 status, reason = self._status_from(r, envl)158 rec = self.remote_gate(node, ws); self.store.add_gate_receipt(run_id, len(gate_receipts) + 1, "controller_post_review", rec); gate_receipts.append(rec)159 # 5. collect artifacts and transcripts160 out_local = os.path.join(rd, "out")161 launcher.rsync_from(node, posixpath.join(ws, "out"), out_local)162 launcher.rsync_from(node, ctl, os.path.join(rd, "ctl"))163 self._collect_transcripts(node, ws, sessions, os.path.join(rd, "sessions"))164 am = provenance.artifact_manifest(out_local)165 json.dump(am, open(os.path.join(rd, "artifact_manifest.json"), "w"), indent=1)166 # 6. common contract check for ALL arms (post hoc; not fed back) + behavioural escape scan167 contract = self.remote_gate(node, ws)168 self.store.add_gate_receipt(run_id, 0, "final_common_contract_check", contract)169 esc = isolation.scan_raw_stream(raw, ws)170 json.dump(esc, open(os.path.join(rd, "behaviour_audit.json"), "w"), indent=1)171 if fault_rec is None and fault and fault.get("type") == "missing_fixture":172 fault_rec = self.store.get_run(run_id).get("fault")173 self.store.transition(run_id, status, note=reason, terminal_reason=reason, ended_at_utc=utcnow(), session_ids=sessions,174 usage=envl.summary(), artifact_manifest=os.path.join("runs", m["split"], run_id, "artifact_manifest.json"),175 fault=fault_rec, model_reported=self._model_reported(run_id))176 # 7. evaluate (laptop side, hidden oracle)177 ev = evaluate_run(self.repo, self.store, run_id, contract_receipt=contract)178 self.log(f"[{run_id}] {task_id} arm {arm} on {node}: {status} ({reason}); cost {envl.spent_usd:.3f} USD; success={ev.get('scientific_success')}")179 return {"run_id": run_id, "status": status, "reason": reason, "cost": envl.spent_usd, "evaluation": ev}180181 def _status_from(self, r, envl):182 if r.get("timed_out"):183 return "timeout", "wall_clock_limit"184 sub = (r.get("result") or {}).get("subtype")185 if sub == "error_max_budget_usd" or (envl.remaining_usd() <= 0.0):186 return "budget_exhausted", sub or "budget_exhausted"187 if sub == "error_max_turns":188 return "completed", "max_turns_reached"189 if r.get("exit_code") not in (0, None) and not r.get("result"):190 return "error", f"exit_code_{r.get('exit_code')}:{(r.get('stderr_tail') or '')[-200:]}"191 return "completed", "executor_finished" if sub == "success" else f"result_{sub}"192193 def _session(self, run_id, node, ws, sid, prompt_file, sys_file, envl, raw, seq, role, kind, resume=False, fault=None):194 cap = min(envl.remaining_usd(), self.budget.max_budget_usd)195 argv = launcher.build_command(ws, ws + ".ctl", prompt_file, sys_file, self.model, sid, cap, self.budget.max_turns_per_session, resume=resume, effort=self.effort)196 job_key = f"{run_id}:{sid}:{kind}"197 if not self.store.log_job(job_key, run_id, "session", json.dumps({"node": node, "argv": argv})):198 raise RuntimeError(f"duplicate session job {job_key}")199 self.store.log_session(sid, run_id, role, kind, seq, {"node": node, "argv": argv, "cap_usd": cap, "wall_s": envl.remaining_wall()}, raw)200 if role == "executor" and resume is False and self.store.query("SELECT 1 FROM sessions WHERE session_id=? AND seq<>?", (sid, seq)):201 pass202 on_event = None; fired = {}203 if fault and fault.get("type") == "interrupt":204 trig = fault.get("trigger", {"after_tool_uses": 6}); counter = {"n": 0}205 def on_event(ev):206 if ev.get("event_type") == "assistant" and ev.get("tool_name"):207 counter["n"] += 1208 if counter["n"] == trig.get("after_tool_uses") and not fired:209 fired.update({"fault_type": "interrupt", "trigger": trig, "at_event": ev["event_id"], "utc": utcnow()})210 launcher.kill_remote(node, sid)211 try: snap = isolation.remote_manifest(node, ws)212 except Exception: snap = {} # noqa: BLE001213 self.store.add_fault(run_id, "interrupt", json.dumps(trig), {"at_event": ev["event_id"], "workspace_manifest_at_kill": snap})214 fired["workspace_manifest_at_kill"] = snap215 r = launcher.run_session(node, ws, argv, raw, run_id, sid, self.store, envl.remaining_wall(), seq_start=seq, on_event=on_event)216 res = r.get("result") or {}217 prev_total = sum(s["cost_usd"] for s in envl.sessions if s["session_id"] == sid)218 total = float(res.get("total_cost_usd") or 0.0)219 session_cost = max(0.0, total - prev_total) if resume else total220 if not res and r.get("timed_out"):221 session_cost = self._estimate_cost_from_events(run_id, sid, seq, r["seq_end"])222 envl.charge(sid, role, session_cost, r["wall_s"], res.get("num_turns"))223 first_in = self.store.query("SELECT usage_json FROM events WHERE run_id=? AND session_id=? AND seq>? AND event_type='assistant' ORDER BY seq LIMIT 1", (run_id, sid, seq))224 fti = None225 if first_in and first_in[0]["usage_json"]:226 u = json.loads(first_in[0]["usage_json"]); fti = (u.get("input_tokens", 0) + u.get("cache_creation_input_tokens", 0) + u.get("cache_read_input_tokens", 0))227 self.store.finish_session(sid, started_utc=r["started_utc"], ended_utc=r["ended_utc"], exit_code=r["exit_code"], result_subtype=res.get("subtype"),228 is_error=int(bool(res.get("is_error"))), total_cost_usd=total, session_cost_usd=session_cost, num_turns=res.get("num_turns"),229 duration_ms=res.get("duration_ms"), usage_json=json.dumps(res.get("usage")), model_usage_json=json.dumps(res.get("modelUsage")),230 first_turn_input_tokens=fti, model_reported=",".join((res.get("modelUsage") or {}).keys()) or None,231 note=("timed_out" if r["timed_out"] else "") + (" fault_fired" if fired else ""))232 self.store.set_job(job_key, "done" if not r["timed_out"] else "timed_out")233 if fired:234 r["fault_fired"] = fired235 return r, r["seq_end"]236237 def _estimate_cost_from_events(self, run_id, sid, seq0, seq1):238 """When a session is killed before its result line, sum per-message usage into a list-price estimate (recorded as such)."""239 rows = self.store.query("SELECT usage_json FROM events WHERE run_id=? AND session_id=? AND seq>? AND seq<=? AND event_type='assistant'", (run_id, sid, seq0, seq1))240 p = self.cfg.get("list_prices_per_mtok", {"input": 4.0, "output": 20.0, "cache_read": 0.20, "cache_write": 5.0})241 tot = 0.0242 for r in rows:243 if not r["usage_json"]: continue244 u = json.loads(r["usage_json"])245 tot += (u.get("input_tokens", 0) * p["input"] + u.get("output_tokens", 0) * p["output"] + u.get("cache_read_input_tokens", 0) * p["cache_read"] + u.get("cache_creation_input_tokens", 0) * p["cache_write"]) / 1e6246 return tot247248 def _model_reported(self, run_id):249 rows = self.store.query("SELECT model_reported FROM sessions WHERE run_id=? AND model_reported IS NOT NULL", (run_id,))250 return ",".join(sorted({r["model_reported"] for r in rows})) or None251252 def remote_gate(self, node, ws) -> dict:253 """Run the pristine gate copy installed on the node (outside the workspace) and return its receipt."""254 r = launcher.ssh(node, f"cd {shlex.quote(ws)} && ~/wp29/venv/bin/python ~/wp29/gates/pgar_gate.py . --python ~/wp29/venv/bin/python --json", timeout=900)255 try:256 return json.loads(r.stdout[r.stdout.index("{"):])257 except Exception: # noqa: BLE001258 return {"passed": False, "n_checks": 0, "n_failed": 1, "checks": [{"id": "gate_execution", "passed": False, "detail": (r.stderr or r.stdout)[-500:]}], "gate_version": "unavailable"}259260 @staticmethod261 def _gate_text(rec) -> str:262 lines = [f"PGAR GATE — {'PASSED' if rec.get('passed') else 'FAILED'} ({rec.get('n_failed')} of {rec.get('n_checks')} checks failed)"]263 for c in rec.get("checks", []):264 if not c.get("passed"):265 lines.append(f" [FAIL] {c['id']}: {c.get('detail')}")266 return "\n".join(lines) + "\n"267268 def _fetch_json(self, node, path):269 r = launcher.ssh(node, f"cat {shlex.quote(path)} 2>/dev/null")270 try: return json.loads(r.stdout)271 except Exception: return None # noqa: BLE001272273 def _collect_transcripts(self, node, ws, session_ids, dst):274 os.makedirs(dst, exist_ok=True)275 for sid in session_ids:276 r = launcher.ssh(node, f"find ~/.claude/projects -name {shlex.quote(sid + '.jsonl')} 2>/dev/null | head -1")277 p = r.stdout.strip()278 if p:279 launcher.rsync_from(node, os.path.dirname(p), dst, excludes=["*"]) if False else None280 import subprocess281 subprocess.run(["scp", "-q", "-o", "BatchMode=yes", f"{node}:{p}", os.path.join(dst, sid + ".jsonl")], capture_output=True, timeout=300)282283 # ---------- batch execution284 def run_many(self, run_ids: list[str], max_workers: int | None = None, faults: dict | None = None) -> list[dict]:285 max_workers = max_workers or len(self.nodes) * int(self.cfg.get("concurrency_per_node", 2))286 out = []287 with cf.ThreadPoolExecutor(max_workers=max_workers) as ex:288 futs = {ex.submit(self.execute, rid, (faults or {}).get(rid)): rid for rid in run_ids}289 for f in cf.as_completed(futs):290 out.append(f.result())291 return out292