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 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 res = launcher.run_session(node, ws, argv, os.path.join(raw_dir, "raw_stream.jsonl"), run_id, session_id, store, wall)53 store.update_run(run_id, notes=f"caps budget={budget} wall={wall}")54 r = res["result"] or {}55 cost = float(r.get("total_cost_usd") or 0.0)56 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"),57 cost_usd=cost, num_turns=r.get("num_turns"), n_events=res["seq_end"], notes=(res["stderr_tail"] or "")[-500:])58 store.ledger_add(run_id, "worker", cost)59 # ---- collect60 col = bundle.collect(node, run_id, raw_dir)61 first_ok = None62 if cfg["recovery"]:63 # local self-check on the collected deliverables (same public tool the worker had)64 tmp = os.path.join(raw_dir, "_sc"); os.makedirs(tmp, exist_ok=True)65 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)66 sc = selfcheck_local(tmp); first_ok = bool(sc["ok"]); store.update_run(run_id, selfcheck_pass_first=int(first_ok))67 ended_badly = res["timed_out"] or (r.get("subtype") not in (None, "success")) or r.get("is_error")68 if (not first_ok or ended_badly) and not STOP.is_set():69 remaining_budget = max(0.5, budget - cost); remaining_wall = max(300, wall - int(res["wall_s"]))70 msg = ("RECOVERY SESSION. Your previous session on this assignment ended" + (" by timeout/budget/error" if ended_badly else "") +71 ". The public self-check reports: " + json.dumps(sc.get("problems", []))[:1500] +72 ". 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.")73 open(os.path.join(raw_dir, "RECOVERY_PROMPT.md"), "w").write(msg)74 launcher.ssh(node, f"cat > {ws}/RECOVERY.md <<'EOF'\n{msg}\nEOF")75 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)76 store.update_run(run_id, status="recovering", recovery_attempted=1)77 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"])78 r2 = res2["result"] or {}79 c2 = float(r2.get("total_cost_usd") or 0.0)80 rec_cost = max(0.0, c2 - cost) if c2 >= cost else c2 # cumulative on resume (WP29 gotcha); fall back if not cumulative81 store.update_run(run_id, recovery_cost_usd=rec_cost, recovery_wall_s=res2["wall_s"]); store.ledger_add(run_id, "recovery", rec_cost)82 shutil.rmtree(os.path.join(raw_dir, "out"), ignore_errors=True)83 col = bundle.collect(node, run_id, raw_dir)84 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)85 sc2 = selfcheck_local(tmp); store.update_run(run_id, selfcheck_pass_final=int(bool(sc2["ok"])), recovery_success=int(bool(sc2["ok"])))86 else:87 store.update_run(run_id, selfcheck_pass_final=int(first_ok))88 shutil.rmtree(tmp, ignore_errors=True)89 store.add_artifacts(run_id, [(p, v["sha256"], v["bytes"]) for p, v in col["artifacts"].items()])90 json.dump(col, open(os.path.join(raw_dir, "collection.json"), "w"), indent=1)91 store.update_run(run_id, status="collected", ended_utc=utcnow())92 except Exception as e: # noqa: BLE00193 store.update_run(run_id, status="infra_error", ended_utc=utcnow(), notes=("infra: " + repr(e) + "\n" + traceback.format_exc())[-1500:])94 log(f"[{run_id}] INFRA ERROR {e!r}")95 finally:96 try:97 if cfg["portal"]: bundle.stop_portal(node, run_id)98 bundle.reset(node, run_id)99 except Exception: pass # noqa: BLE001100101def evaluate_and_judge(store, run_id: str, log=print, do_judge=True):102 run = store.get_run(run_id); raw_dir = os.path.join(RAW_RUNS, run_id); out_dir = os.path.join(raw_dir, "out")103 col = json.load(open(os.path.join(raw_dir, "collection.json"))) if os.path.exists(os.path.join(raw_dir, "collection.json")) else {}104 receipt = ev_mod.evaluate_out(run["task_id"], out_dir, run_id, col.get("inputs_changed"))105 json.dump(receipt, open(os.path.join(raw_dir, "receipt.json"), "w"), indent=1, default=str); store.set_receipt(run_id, receipt)106 if do_judge:107 js = judge_mod.judge_run(run["task_id"], out_dir, store, run_id)108 json.dump(js, open(os.path.join(raw_dir, "judgments.json"), "w"), indent=1)109 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,)))110 try:111 comp = json.load(open(os.path.join(out_dir, "COMPLETION.json")))112 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 [])113 except Exception: # noqa: BLE001114 m["completion_status"] = None; m["n_questions"] = None; m["n_issues_found"] = None; m["n_assumptions"] = None115 m["n_artifacts"] = len(col.get("artifacts") or {}); m["inputs_changed"] = bool(col.get("inputs_changed"))116 store.set_metrics(run_id, m); store.update_run(run_id, status="evaluated")117 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']}")118 return m119120def run_queue(store, where="status='assigned'", max_parallel_total=12, log=print, do_judge=True, envelope=STUDY_ENVELOPE_USD):121 busy = {n: 0 for n in NODES}; lock = threading.Lock(); threads = []122 pending = [r["run_id"] for r in store.runs(where)]123 log(f"queue: {len(pending)} runs; ledger so far {store.ledger_total():.2f} USD")124 def worker(rid, node):125 try:126 execute_run(store, rid, node, log)127 if store.get_run(rid)["status"] == "collected": evaluate_and_judge(store, rid, log, do_judge)128 finally:129 with lock: busy[node] -= 1130 while pending and not STOP.is_set():131 if store.ledger_total() >= envelope:132 log(f"ENVELOPE REACHED ({store.ledger_total():.2f} >= {envelope}); stopping"); break133 rid = pending[0]; arm = store.get_run(rid)["arm"]134 with lock:135 node = pick_node(arm, busy) if sum(busy.values()) < max_parallel_total else None136 if node: busy[node] += 1137 if not node: time.sleep(5); continue138 pending.pop(0); t = threading.Thread(target=worker, args=(rid, node), daemon=True); t.start(); threads.append(t); time.sleep(2)139 for t in threads: t.join()140 log(f"queue done; ledger {store.ledger_total():.2f} USD")141