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="", interrupt_s=None):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, interrupt_s=interrupt_s); 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 task = bundle.load_task(task_id)38 budget = max(MAX_BUDGET_USD, float(task.get("budget_usd") or 0)); wall = max(WALL_SECONDS, int(task.get("time_limit_s") or 0)) # per-scenario caps (A-006)39 raw_dir = os.path.join(RAW_RUNS, run_id); os.makedirs(raw_dir, exist_ok=True)40 session_id = new_session_id(); portal_url = None; port = None41 store.update_run(run_id, node=node, session_id=session_id, status="preparing", started_utc=utcnow())42 try:43 if cfg["portal"]:44 port = free_port(node); portal_url = f"http://127.0.0.1:{port}"45 pj = os.path.join(bundle.scenario_dir(task_id), "portal", "portal.json")46 puser = json.load(open(pj)).get("user") if (cfg["portal"] and os.path.exists(pj)) else None47 local_ws = bundle.compose_local(task_id, arm, run_id, portal_url, puser)48 ws = bundle.push(node, local_ws, run_id); shutil.rmtree(local_ws, ignore_errors=True)49 if cfg["portal"]: bundle.start_portal(node, run_id, port)50 argv = launcher.build_argv(f"{ws}/PROMPT.md", f"{ws}/SYSTEM.md", run["model"], session_id, budget, MAX_TURNS, cfg["tools"], cfg["disallowed"])51 store.update_run(run_id, status="running")52 first_wall = int(run["interrupt_s"]) if run.get("interrupt_s") else wall # recovery experiment: the first session is cut at interrupt_s53 res = launcher.run_session(node, ws, argv, os.path.join(raw_dir, "raw_stream.jsonl"), run_id, session_id, store, first_wall)54 store.update_run(run_id, notes=f"caps budget={budget} wall={wall}")55 r = res["result"] or {}56 cost = float(r.get("total_cost_usd") or 0.0)57 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"),58 cost_usd=cost, num_turns=r.get("num_turns"), n_events=res["seq_end"], notes=(res["stderr_tail"] or "")[-500:])59 store.ledger_add(run_id, "worker", cost)60 # ---- collect61 col = bundle.collect(node, run_id, raw_dir)62 first_ok = None63 if cfg["recovery"]:64 # local self-check on the collected deliverables (same public tool the worker had)65 tmp = os.path.join(raw_dir, "_sc"); os.makedirs(tmp, exist_ok=True)66 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)67 sc = selfcheck_local(tmp); first_ok = bool(sc["ok"]); store.update_run(run_id, selfcheck_pass_first=int(first_ok))68 ended_badly = res["timed_out"] or (r.get("subtype") not in (None, "success")) or r.get("is_error")69 if (not first_ok or ended_badly) and not STOP.is_set():70 remaining_budget = max(0.5, budget - cost); remaining_wall = max(300, wall - int(res["wall_s"])) # interrupted runs keep the scenario wall minus what was used71 msg = ("RECOVERY SESSION. Your previous session on this assignment ended" + (" by timeout/budget/error" if ended_badly else "") +72 ". The public self-check reports: " + json.dumps(sc.get("problems", []))[:1500] +73 ". 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.")74 open(os.path.join(raw_dir, "RECOVERY_PROMPT.md"), "w").write(msg)75 launcher.ssh(node, f"cat > {ws}/RECOVERY.md <<'EOF'\n{msg}\nEOF")76 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)77 store.update_run(run_id, status="recovering", recovery_attempted=1)78 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"])79 r2 = res2["result"] or {}80 c2 = float(r2.get("total_cost_usd") or 0.0)81 rec_cost = max(0.0, c2 - cost) if c2 >= cost else c2 # cumulative on resume (WP29 gotcha); fall back if not cumulative82 store.update_run(run_id, recovery_cost_usd=rec_cost, recovery_wall_s=res2["wall_s"]); store.ledger_add(run_id, "recovery", rec_cost)83 shutil.rmtree(os.path.join(raw_dir, "out"), ignore_errors=True)84 col = bundle.collect(node, run_id, raw_dir)85 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)86 sc2 = selfcheck_local(tmp); store.update_run(run_id, selfcheck_pass_final=int(bool(sc2["ok"])), recovery_success=int(bool(sc2["ok"])))87 else:88 store.update_run(run_id, selfcheck_pass_final=int(first_ok))89 shutil.rmtree(tmp, ignore_errors=True)90 store.add_artifacts(run_id, [(p, v["sha256"], v["bytes"]) for p, v in col["artifacts"].items()])91 json.dump(col, open(os.path.join(raw_dir, "collection.json"), "w"), indent=1)92 store.update_run(run_id, status="collected", ended_utc=utcnow())93 except Exception as e: # noqa: BLE00194 store.update_run(run_id, status="infra_error", ended_utc=utcnow(), notes=("infra: " + repr(e) + "\n" + traceback.format_exc())[-1500:])95 log(f"[{run_id}] INFRA ERROR {e!r}")96 finally:97 try:98 if cfg["portal"]: bundle.stop_portal(node, run_id)99 bundle.reset(node, run_id)100 except Exception: pass # noqa: BLE001101102def evaluate_and_judge(store, run_id: str, log=print, do_judge=True):103 run = store.get_run(run_id); raw_dir = os.path.join(RAW_RUNS, run_id); out_dir = os.path.join(raw_dir, "out")104 col = json.load(open(os.path.join(raw_dir, "collection.json"))) if os.path.exists(os.path.join(raw_dir, "collection.json")) else {}105 receipt = ev_mod.evaluate_out(run["task_id"], out_dir, run_id, col.get("inputs_changed"))106 json.dump(receipt, open(os.path.join(raw_dir, "receipt.json"), "w"), indent=1, default=str); store.set_receipt(run_id, receipt)107 if do_judge:108 js = judge_mod.judge_run(run["task_id"], out_dir, store, run_id)109 json.dump(js, open(os.path.join(raw_dir, "judgments.json"), "w"), indent=1)110 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,)))111 try:112 comp = json.load(open(os.path.join(out_dir, "COMPLETION.json")))113 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 [])114 except Exception: # noqa: BLE001115 m["completion_status"] = None; m["n_questions"] = None; m["n_issues_found"] = None; m["n_assumptions"] = None116 m["n_artifacts"] = len(col.get("artifacts") or {}); m["inputs_changed"] = bool(col.get("inputs_changed"))117 store.set_metrics(run_id, m); store.update_run(run_id, status="evaluated")118 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']}")119 return m120121def run_queue(store, where="status='assigned'", max_parallel_total=12, log=print, do_judge=True, envelope=STUDY_ENVELOPE_USD):122 busy = {n: 0 for n in NODES}; lock = threading.Lock(); threads = []123 pending = [r["run_id"] for r in store.runs(where)]124 log(f"queue: {len(pending)} runs; ledger so far {store.ledger_total():.2f} USD")125 def worker(rid, node):126 try:127 execute_run(store, rid, node, log)128 if store.get_run(rid)["status"] == "collected": evaluate_and_judge(store, rid, log, do_judge)129 finally:130 with lock: busy[node] -= 1131 while pending and not STOP.is_set():132 if store.ledger_total() >= envelope:133 log(f"ENVELOPE REACHED ({store.ledger_total():.2f} >= {envelope}); stopping"); break134 rid = pending[0]; arm = store.get_run(rid)["arm"]135 with lock:136 node = pick_node(arm, busy) if sum(busy.values()) < max_parallel_total else None137 if node: busy[node] += 1138 if not node: time.sleep(5); continue139 pending.pop(0); t = threading.Thread(target=worker, args=(rid, node), daemon=True); t.start(); threads.append(t); time.sleep(2)140 for t in threads: t.join()141 log(f"queue done; ledger {store.ledger_total():.2f} USD")142