1"""Assignment (seeded randomisation) and execution pipeline with per-node concurrency; one bounded recovery for arm A4."""2from __future__ import annotations3import json, os, random, shutil, subprocess, sys, threading, time, traceback4from .config import (REPO, RAW_RUNS, NODES, PRIMARY_MODEL, MAX_BUDGET_USD, WALL_SECONDS, MAX_TURNS, STUDY_ENVELOPE_USD, NODE_PYTHON, NODE_HOME)5from .contracts import utcnow, new_session_id, new_run_id6from .arms import ARMS7from . import bundle, launcher, evaluate as ev_mod, judge as judge_mod, metrics as metrics_mod89STOP = threading.Event()1011def make_assignment(store, task_ids, arms, reps, seed, model=PRIMARY_MODEL, stage="pilot", model_tag=""):12 rng = random.Random(seed); cells = [(t, a, r) for t in task_ids for a in arms for r in range(1, reps + 1)]13 rng.shuffle(cells); n = 014 for i, (t, a, r) in enumerate(cells):15 task = bundle.load_task(t)16 rid = new_run_id(t, a, r, model_tag)17 if store.get_run(rid): continue18 store.add_run(run_id=rid, task_id=t, occupation=task.get("occupation"), arm=a, rep=r, model=model, stage=stage, order_index=i, seed=seed); n += 119 return n2021def pick_node(arm: str, busy: dict) -> str | None:22 cands = [n for n, c in NODES.items() if (c["system_events"] or not ARMS[arm]["requires_system_events"]) and busy[n] < c["slots"]]23 if not cands: return None24 return min(cands, key=lambda n: busy[n] / NODES[n]["slots"])2526def free_port(node: str) -> int:27 r = launcher.ssh(node, f"{NODE_PYTHON} -c 'import socket;s=socket.socket();s.bind((\"127.0.0.1\",0));print(s.getsockname()[1])'")28 return int(r.stdout.strip())2930def selfcheck_local(ws_dir: str) -> dict:31 r = subprocess.run([sys.executable, os.path.join(REPO, "benchmark", "tools", "selfcheck.py"), "--workspace", ws_dir, "--json"], capture_output=True, text=True)32 try: return json.loads(r.stdout)33 except Exception: return {"ok": False, "problems": ["selfcheck crashed: " + r.stderr[-300:]]} # noqa: BLE0013435def execute_run(store, run_id: str, node: str, log=print):36 run = store.get_run(run_id); arm = run["arm"]; task_id = run["task_id"]; cfg = ARMS[arm]37 raw_dir = os.path.join(RAW_RUNS, run_id); os.makedirs(raw_dir, exist_ok=True)38 session_id = new_session_id(); portal_url = None; port = None39 store.update_run(run_id, node=node, session_id=session_id, status="preparing", started_utc=utcnow())40 try:41 if cfg["portal"]:42 port = free_port(node); portal_url = f"http://127.0.0.1:{port}"43 pj = os.path.join(bundle.scenario_dir(task_id), "portal", "portal.json")44 puser = json.load(open(pj)).get("user") if (cfg["portal"] and os.path.exists(pj)) else None45 local_ws = bundle.compose_local(task_id, arm, run_id, portal_url, puser)46 ws = bundle.push(node, local_ws, run_id); shutil.rmtree(local_ws, ignore_errors=True)47 if cfg["portal"]: bundle.start_portal(node, run_id, port)48 argv = launcher.build_argv(f"{ws}/PROMPT.md", f"{ws}/SYSTEM.md", run["model"], session_id, MAX_BUDGET_USD, MAX_TURNS, cfg["tools"], cfg["disallowed"])49 store.update_run(run_id, status="running")50 res = launcher.run_session(node, ws, argv, os.path.join(raw_dir, "raw_stream.jsonl"), run_id, session_id, store, WALL_SECONDS)51 r = res["result"] or {}52 cost = float(r.get("total_cost_usd") or 0.0)53 store.update_run(run_id, wall_s=res["wall_s"], timed_out=int(res["timed_out"]), exit_code=res["exit_code"], result_subtype=r.get("subtype") or ("timeout" if res["timed_out"] else "no_result"),54 cost_usd=cost, num_turns=r.get("num_turns"), n_events=res["seq_end"], notes=(res["stderr_tail"] or "")[-500:])55 store.ledger_add(run_id, "worker", cost)56 # ---- collect57 col = bundle.collect(node, run_id, raw_dir)58 first_ok = None59 if cfg["recovery"]:60 # local self-check on the collected deliverables (same public tool the worker had)61 tmp = os.path.join(raw_dir, "_sc"); os.makedirs(tmp, exist_ok=True)62 shutil.copy2(os.path.join(bundle.scenario_dir(task_id), "task.json"), tmp); shutil.copytree(os.path.join(raw_dir, "out"), os.path.join(tmp, "out"), dirs_exist_ok=True)63 sc = selfcheck_local(tmp); first_ok = bool(sc["ok"]); store.update_run(run_id, selfcheck_pass_first=int(first_ok))64 ended_badly = res["timed_out"] or (r.get("subtype") not in (None, "success")) or r.get("is_error")65 if (not first_ok or ended_badly) and not STOP.is_set():66 remaining_budget = max(0.5, MAX_BUDGET_USD - cost); remaining_wall = max(300, WALL_SECONDS - int(res["wall_s"]))67 msg = ("RECOVERY SESSION. Your previous session on this assignment ended" + (" by timeout/budget/error" if ended_badly else "") +68 ". The public self-check reports: " + json.dumps(sc.get("problems", []))[:1500] +69 ". Read out/CHECKPOINT.md (if present) and the current contents of out/, then finish the remaining deliverables and fix the reported problems. Run `python tools/selfcheck.py` before finishing. Rewrite out/COMPLETION.json at the end.")70 open(os.path.join(raw_dir, "RECOVERY_PROMPT.md"), "w").write(msg)71 launcher.ssh(node, f"cat > {ws}/RECOVERY.md <<'EOF'\n{msg}\nEOF")72 argv2 = launcher.build_argv(f"{ws}/RECOVERY.md", f"{ws}/SYSTEM.md", run["model"], session_id, remaining_budget, MAX_TURNS, cfg["tools"], cfg["disallowed"], resume=True)73 store.update_run(run_id, status="recovering", recovery_attempted=1)74 res2 = launcher.run_session(node, ws, argv2, os.path.join(raw_dir, "raw_stream_recovery.jsonl"), run_id, session_id, store, remaining_wall, seq_start=res["seq_end"])75 r2 = res2["result"] or {}76 c2 = float(r2.get("total_cost_usd") or 0.0)77 rec_cost = max(0.0, c2 - cost) if c2 >= cost else c2 # cumulative on resume (WP29 gotcha); fall back if not cumulative78 store.update_run(run_id, recovery_cost_usd=rec_cost, recovery_wall_s=res2["wall_s"]); store.ledger_add(run_id, "recovery", rec_cost)79 shutil.rmtree(os.path.join(raw_dir, "out"), ignore_errors=True)80 col = bundle.collect(node, run_id, raw_dir)81 shutil.rmtree(os.path.join(tmp, "out"), ignore_errors=True); shutil.copytree(os.path.join(raw_dir, "out"), os.path.join(tmp, "out"), dirs_exist_ok=True)82 sc2 = selfcheck_local(tmp); store.update_run(run_id, selfcheck_pass_final=int(bool(sc2["ok"])), recovery_success=int(bool(sc2["ok"])))83 else:84 store.update_run(run_id, selfcheck_pass_final=int(first_ok))85 shutil.rmtree(tmp, ignore_errors=True)86 store.add_artifacts(run_id, [(p, v["sha256"], v["bytes"]) for p, v in col["artifacts"].items()])87 json.dump(col, open(os.path.join(raw_dir, "collection.json"), "w"), indent=1)88 store.update_run(run_id, status="collected", ended_utc=utcnow())89 except Exception as e: # noqa: BLE00190 store.update_run(run_id, status="infra_error", ended_utc=utcnow(), notes=("infra: " + repr(e) + "\n" + traceback.format_exc())[-1500:])91 log(f"[{run_id}] INFRA ERROR {e!r}")92 finally:93 try:94 if cfg["portal"]: bundle.stop_portal(node, run_id)95 bundle.reset(node, run_id)96 except Exception: pass # noqa: BLE0019798def evaluate_and_judge(store, run_id: str, log=print, do_judge=True):99 run = store.get_run(run_id); raw_dir = os.path.join(RAW_RUNS, run_id); out_dir = os.path.join(raw_dir, "out")100 col = json.load(open(os.path.join(raw_dir, "collection.json"))) if os.path.exists(os.path.join(raw_dir, "collection.json")) else {}101 receipt = ev_mod.evaluate_out(run["task_id"], out_dir, run_id, col.get("inputs_changed"))102 json.dump(receipt, open(os.path.join(raw_dir, "receipt.json"), "w"), indent=1, default=str); store.set_receipt(run_id, receipt)103 if do_judge:104 js = judge_mod.judge_run(run["task_id"], out_dir, store, run_id)105 json.dump(js, open(os.path.join(raw_dir, "judgments.json"), "w"), indent=1)106 m = metrics_mod.compute(store.events(run_id), store.get_run(run_id), receipt, store.q("SELECT * FROM judgments WHERE run_id=? ORDER BY judge_idx", (run_id,)))107 try:108 comp = json.load(open(os.path.join(out_dir, "COMPLETION.json")))109 m["completion_status"] = comp.get("status"); m["n_questions"] = len(comp.get("questions") or []); m["n_issues_found"] = len(comp.get("issues_found") or []); m["n_assumptions"] = len(comp.get("assumptions") or [])110 except Exception: # noqa: BLE001111 m["completion_status"] = None; m["n_questions"] = None; m["n_issues_found"] = None; m["n_assumptions"] = None112 m["n_artifacts"] = len(col.get("artifacts") or {}); m["inputs_changed"] = bool(col.get("inputs_changed"))113 store.set_metrics(run_id, m); store.update_run(run_id, status="evaluated")114 log(f"[{run_id}] det={m['det_completion']} score={m['det_score']} viol={m['n_critical_violations']} quality={m['quality']} cost={m['total_cost_usd']}")115 return m116117def run_queue(store, where="status='assigned'", max_parallel_total=12, log=print, do_judge=True, envelope=STUDY_ENVELOPE_USD):118 busy = {n: 0 for n in NODES}; lock = threading.Lock(); threads = []119 pending = [r["run_id"] for r in store.runs(where)]120 log(f"queue: {len(pending)} runs; ledger so far {store.ledger_total():.2f} USD")121 def worker(rid, node):122 try:123 execute_run(store, rid, node, log)124 if store.get_run(rid)["status"] == "collected": evaluate_and_judge(store, rid, log, do_judge)125 finally:126 with lock: busy[node] -= 1127 while pending and not STOP.is_set():128 if store.ledger_total() >= envelope:129 log(f"ENVELOPE REACHED ({store.ledger_total():.2f} >= {envelope}); stopping"); break130 rid = pending[0]; arm = store.get_run(rid)["arm"]131 with lock:132 node = pick_node(arm, busy) if sum(busy.values()) < max_parallel_total else None133 if node: busy[node] += 1134 if not node: time.sleep(5); continue135 pending.pop(0); t = threading.Thread(target=worker, args=(rid, node), daemon=True); t.start(); threads.append(t); time.sleep(2)136 for t in threads: t.join()137 log(f"queue done; ledger {store.ledger_total():.2f} USD")138