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 self._attempt = attempt80 if not self.store.log_job(f"{run_id}:execute:attempt{attempt}", run_id, "execute", node):81 # duplicate submission: idempotency guard82 self.log(f"[{run_id}] duplicate execute request ignored (job already logged)")83 return {"run_id": run_id, "status": "duplicate_ignored"}84 # 1. prepare workspace85 bundle = self._bundle_dir(task_id)86 prep = isolation.prepare_workspace(node, self.remote_root, run_id, bundle)87 if not prep["verified_identical"]:88 raise RuntimeError(f"workspace manifest mismatch: extra={prep['remote_extra']} missing={prep['remote_missing']}")89 ws, ctl = prep["workspace"], prep["ctl_dir"]90 task_md = self._read("tasks", "public_bundles", task_id, "TASK.md")91 common = self._read("protocol", "arms", "common_preamble.md"); arm_text = self._read("protocol", "arms", f"arm_{arm}.md")92 task_prompt = self._read("protocol", "arms", "task_prompt.md")93 sys_append, ctx_hash = isolation.effective_context(common, arm_text, task_md, task_prompt)94 isolation.write_ctl_files(node, ctl, {"system_append.md": sys_append, "prompt.txt": task_prompt,95 "reviewer_system.md": self._read("protocol", "arms", "reviewer.md"),96 "reviewer_prompt.txt": "Review the deliverables in this workspace now, as instructed, and write out/review/review.json.\n"})97 env = provenance.node_environment(node, launcher.ssh)98 self.store.update_manifest(run_id, node=node, task_hash=sha256_file(os.path.join(bundle, "bundle.json")),99 protocol_hash=sha256_text(sys_append), effective_context_hash=ctx_hash,100 environment_id=env["environment_id"], client_version=env.get("claude"), model_requested=self.model,101 code_commit=provenance.code_commit(self.repo))102 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)103 open(os.path.join(rd, "system_append.md"), "w").write(sys_append)104 self.store.transition(run_id, "prepared", note=f"workspace {ws} on {node}")105 if fault and fault.get("type") == "missing_fixture":106 from . import recovery107 fr = recovery.remove_fixture(node, ws, fault["path"])108 self.store.add_fault(run_id, "missing_fixture", fault["path"], fr)109 self.store.update_manifest(run_id, fault={"fault_type": "missing_fixture", "path": fault["path"], "utc": fr["utc"], "removed": fr["ok"]})110 # 2. executor session(s)111 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)112 raw = os.path.join(rd, "raw_stream.jsonl")113 sid = new_session_id(); seq = 0; sessions = []; status = "completed"; reason = "executor_finished"114 self.store.transition(run_id, "running", note="first session", started_at_utc=utcnow(), session_ids=[sid])115 prompt_file = posixpath.join(ctl, "prompt.txt"); sys_file = posixpath.join(ctl, "system_append.md")116 fault_rec = None117 r, seq = self._session(run_id, node, ws, sid, prompt_file, sys_file, envl, raw, seq, role="executor", kind="initial", fault=fault)118 sessions.append(sid)119 if r.get("fault_fired"):120 fault_rec = r["fault_fired"]121 # standardized interruption → resume once from the same session122 self.store.transition(run_id, "interrupted", note="fault: " + fault_rec["fault_type"])123 if envl.can_continue():124 isolation.write_ctl_files(node, ctl, {"resume_prompt.txt": self._read("protocol", "arms", "resume_after_interruption_prompt.md")})125 self.store.transition(run_id, "running", note="resume after interruption")126 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)127 status, reason = self._status_from(r, envl)128 # 3. controller-enforced gate cycles for C/D129 gate_receipts = []130 if arm in ("C", "D") and status == "completed":131 for cycle in range(1, self.budget.max_gate_cycles + 1):132 rec = self.remote_gate(node, ws)133 self.store.add_gate_receipt(run_id, cycle, "controller", rec); gate_receipts.append(rec)134 if rec.get("passed"):135 break136 if not envl.can_continue():137 reason = "gate_failed_budget_exhausted"; break138 fb = self._read("protocol", "arms", "gate_feedback_prompt.md") + self._gate_text(rec)139 isolation.write_ctl_files(node, ctl, {f"gate_feedback_{cycle}.txt": fb})140 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)141 status, reason = self._status_from(r, envl)142 if status != "completed":143 break144 else:145 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)146 # 4. separate review for D147 review = None148 if arm == "D" and status == "completed" and gate_receipts and gate_receipts[-1].get("passed") and envl.can_continue():149 rsid = new_session_id(); sessions.append(rsid)150 launcher.ssh(node, f"mkdir -p {shlex.quote(ws)}/out/review")151 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")152 review = self._fetch_json(node, posixpath.join(ws, "out", "review", "review.json"))153 if review is None: # reviewer sessions can drift their cwd (observed: out/out/review/review.json); search under out/154 r_ = launcher.ssh(node, f"find {shlex.quote(posixpath.join(ws, 'out'))} -name review.json | head -1"); alt = r_.stdout.strip()155 if alt: review = self._fetch_json(node, alt)156 self.store.add_review(run_id, rsid, review)157 if review and review.get("verdict") == "revise" and envl.can_continue():158 fb = self._read("protocol", "arms", "review_feedback_prompt.md") + json.dumps(review, indent=1)[:6000]159 isolation.write_ctl_files(node, ctl, {"review_feedback.txt": fb})160 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)161 status, reason = self._status_from(r, envl)162 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)163 # 5. collect artifacts and transcripts164 out_local = os.path.join(rd, "out")165 launcher.rsync_from(node, posixpath.join(ws, "out"), out_local)166 launcher.rsync_from(node, ctl, os.path.join(rd, "ctl"))167 self._collect_transcripts(node, ws, sessions, os.path.join(rd, "sessions"))168 am = provenance.artifact_manifest(out_local)169 json.dump(am, open(os.path.join(rd, "artifact_manifest.json"), "w"), indent=1)170 # 6. common contract check for ALL arms (post hoc; not fed back) + behavioural escape scan171 contract = self.remote_gate(node, ws)172 self.store.add_gate_receipt(run_id, 0, "final_common_contract_check", contract)173 esc = isolation.scan_raw_stream(raw, ws)174 json.dump(esc, open(os.path.join(rd, "behaviour_audit.json"), "w"), indent=1)175 if fault_rec is None and fault and fault.get("type") == "missing_fixture":176 fault_rec = self.store.get_run(run_id).get("fault")177 self.store.transition(run_id, status, note=reason, terminal_reason=reason, ended_at_utc=utcnow(), session_ids=sessions,178 usage=envl.summary(), artifact_manifest=os.path.join("runs", m["split"], run_id, "artifact_manifest.json"),179 fault=fault_rec, model_reported=self._model_reported(run_id))180 # 7. evaluate (laptop side, hidden oracle)181 ev = evaluate_run(self.repo, self.store, run_id, contract_receipt=contract)182 self.log(f"[{run_id}] {task_id} arm {arm} on {node}: {status} ({reason}); cost {envl.spent_usd:.3f} USD; success={ev.get('scientific_success')}")183 return {"run_id": run_id, "status": status, "reason": reason, "cost": envl.spent_usd, "evaluation": ev}184185 def _status_from(self, r, envl):186 if r.get("timed_out"):187 return "timeout", "wall_clock_limit"188 sub = (r.get("result") or {}).get("subtype")189 if sub == "error_max_budget_usd" or (envl.remaining_usd() <= 0.0):190 return "budget_exhausted", sub or "budget_exhausted"191 if sub == "error_max_turns":192 return "completed", "max_turns_reached"193 if r.get("exit_code") not in (0, None) and not r.get("result"):194 return "error", f"exit_code_{r.get('exit_code')}:{(r.get('stderr_tail') or '')[-200:]}"195 return "completed", "executor_finished" if sub == "success" else f"result_{sub}"196197 def _session(self, run_id, node, ws, sid, prompt_file, sys_file, envl, raw, seq, role, kind, resume=False, fault=None):198 cap = min(envl.remaining_usd(), self.budget.max_budget_usd)199 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)200 job_key = f"{run_id}:{sid}:{kind}"201 if not self.store.log_job(job_key, run_id, "session", json.dumps({"node": node, "argv": argv})):202 raise RuntimeError(f"duplicate session job {job_key}")203 row_key = sid if not resume else f"{sid}#{seq}" # continuation sessions share the client session id; the ledger row is keyed per launch204 self.store.log_session(row_key, run_id, role, kind, seq, {"node": node, "argv": argv, "cap_usd": cap, "wall_s": envl.remaining_wall(), "client_session_id": sid}, raw)205 on_event = None; fired = {}206 if fault and fault.get("type") == "interrupt":207 trig = fault.get("trigger", {"after_tool_uses": 6}); counter = {"n": 0}208 def on_event(ev):209 if ev.get("event_type") == "assistant" and ev.get("tool_name"):210 counter["n"] += 1211 if counter["n"] == trig.get("after_tool_uses") and not fired:212 fired.update({"fault_type": "interrupt", "trigger": trig, "at_event": ev["event_id"], "utc": utcnow()})213 launcher.kill_remote(node, sid)214 try: snap = isolation.remote_manifest(node, ws)215 except Exception: snap = {} # noqa: BLE001216 self.store.add_fault(run_id, "interrupt", json.dumps(trig), {"at_event": ev["event_id"], "workspace_manifest_at_kill": snap})217 fired["workspace_manifest_at_kill"] = snap218 r = launcher.run_session(node, ws, argv, raw, run_id, sid, self.store, envl.remaining_wall(), seq_start=seq, on_event=on_event, event_prefix=(f"{run_id}:a{self._attempt}" if getattr(self, "_attempt", 0) else run_id))219 res = r.get("result") or {}220 prev_total = sum(s["cost_usd"] for s in envl.sessions if s["session_id"] == sid)221 total = float(res.get("total_cost_usd") or 0.0)222 session_cost = max(0.0, total - prev_total) if resume else total223 if not res and r.get("timed_out"):224 session_cost = self._estimate_cost_from_events(run_id, sid, seq, r["seq_end"])225 envl.charge(sid, role, session_cost, r["wall_s"], res.get("num_turns"))226 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))227 fti = None228 if first_in and first_in[0]["usage_json"]:229 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))230 self.store.finish_session(row_key, started_utc=r["started_utc"], ended_utc=r["ended_utc"], exit_code=r["exit_code"], result_subtype=res.get("subtype"),231 is_error=int(bool(res.get("is_error"))), total_cost_usd=total, session_cost_usd=session_cost, num_turns=res.get("num_turns"),232 duration_ms=res.get("duration_ms"), usage_json=json.dumps(res.get("usage")), model_usage_json=json.dumps(res.get("modelUsage")),233 first_turn_input_tokens=fti, model_reported=",".join((res.get("modelUsage") or {}).keys()) or None,234 note=("timed_out" if r["timed_out"] else "") + (" fault_fired" if fired else ""))235 self.store.set_job(job_key, "done" if not r["timed_out"] else "timed_out")236 if fired:237 r["fault_fired"] = fired238 return r, r["seq_end"]239240 def _estimate_cost_from_events(self, run_id, sid, seq0, seq1):241 """When a session is killed before its result line, sum per-message usage into a list-price estimate (recorded as such)."""242 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))243 p = self.cfg.get("list_prices_per_mtok", {"input": 4.0, "output": 20.0, "cache_read": 0.20, "cache_write": 5.0})244 tot = 0.0245 for r in rows:246 if not r["usage_json"]: continue247 u = json.loads(r["usage_json"])248 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"]) / 1e6249 return tot250251 def _model_reported(self, run_id):252 rows = self.store.query("SELECT model_reported FROM sessions WHERE run_id=? AND model_reported IS NOT NULL", (run_id,))253 return ",".join(sorted({r["model_reported"] for r in rows})) or None254255 def remote_gate(self, node, ws) -> dict:256 """Run the pristine gate copy installed on the node (outside the workspace) and return its receipt."""257 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)258 try:259 return json.loads(r.stdout[r.stdout.index("{"):])260 except Exception: # noqa: BLE001261 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"}262263 @staticmethod264 def _gate_text(rec) -> str:265 lines = [f"PGAR GATE — {'PASSED' if rec.get('passed') else 'FAILED'} ({rec.get('n_failed')} of {rec.get('n_checks')} checks failed)"]266 for c in rec.get("checks", []):267 if not c.get("passed"):268 lines.append(f" [FAIL] {c['id']}: {c.get('detail')}")269 return "\n".join(lines) + "\n"270271 def _fetch_json(self, node, path):272 r = launcher.ssh(node, f"cat {shlex.quote(path)} 2>/dev/null")273 try: return json.loads(r.stdout)274 except Exception: return None # noqa: BLE001275276 def _collect_transcripts(self, node, ws, session_ids, dst):277 os.makedirs(dst, exist_ok=True)278 for sid in session_ids:279 r = launcher.ssh(node, f"find ~/.claude/projects -name {shlex.quote(sid + '.jsonl')} 2>/dev/null | head -1")280 p = r.stdout.strip()281 if p:282 import subprocess283 subprocess.run(["scp", "-q", "-o", "BatchMode=yes", f"{node}:{p}", os.path.join(dst, sid + ".jsonl")], capture_output=True, timeout=300)284285 # ---------- batch execution286 def run_many(self, run_ids: list[str], max_workers: int | None = None, faults: dict | None = None) -> list[dict]:287 max_workers = max_workers or len(self.nodes) * int(self.cfg.get("concurrency_per_node", 2))288 out = []289 with cf.ThreadPoolExecutor(max_workers=max_workers) as ex:290 futs = {ex.submit(self.execute, rid, (faults or {}).get(rid)): rid for rid in run_ids}291 for f in cf.as_completed(futs):292 out.append(f.result())293 return out294