#!/usr/bin/env python3 """Science controller: validate Pi, run fixed-budget episodes, grade saved artifacts.""" from __future__ import annotations import argparse from datetime import datetime, timezone import hashlib import json import math import os from pathlib import Path import signal import subprocess import sys import threading import time from budget_proxy import Gateway import programbench_adapter as benchmark HERE = Path(__file__).resolve().parent def save(path, value): path = Path(path) path.parent.mkdir(parents=True, exist_ok=True) temp = path.with_suffix(path.suffix + ".tmp") temp.write_text(json.dumps(value, indent=2) + "\n") temp.replace(path) def load(path): return json.loads(Path(path).read_text()) def safe_env(local_key=None): # Docker subprocesses inherit routing, but neither hosted-service credential. env = {k:v for k,v in os.environ.items() if k not in {"HF_TOKEN","HUGGING_FACE_HUB_TOKEN","REFLECTION_API_KEY"}} if local_key: env["PI_STUDY_API_KEY"] = local_key return env class Study: def __init__(self, args): self.args = args self.root = args.output_dir.resolve() self.root.mkdir(parents=True, exist_ok=True) self.config = load(args.config) self.manifest = benchmark.load_manifest(args.manifest) self.events = self.root / "events.jsonl" self.cancelled = False self.current = None self.current_is_adapter = False self.event_lock = threading.Lock() self.local_key = (args.secrets_dir / "proxy-key").read_text().strip() self.gateway = Gateway(("127.0.0.1", 0), self.config, self.root / "budget.sqlite", self.events, (args.secrets_dir / "reflection-api-key").read_text().strip(),self.local_key) self.gateway.event_lock = self.event_lock self.server_thread = threading.Thread(target=self.gateway.serve_forever,daemon=True) self.server_thread.start() tracking_env = safe_env() tracking_env["HF_TOKEN"] = (args.secrets_dir / "hf-token").read_text().strip() self.tracking_log = (self.root / "tracking.log").open("a") self.tracker = subprocess.Popen([sys.executable,str(HERE / "tracking.py"),"--events",str(self.events), "--state-dir",str(self.root / "tracking"),"--space-id",self.config["space_id"],"--project",self.config["project"]], env=tracking_env,stdout=self.tracking_log,stderr=subprocess.STDOUT) self.started = time.monotonic() self.emit("study-summary","run_started",status="running",config={ "model":self.config["model"],"phase":"study_summary","tasks":5,"conditions":3,"repetitions":1, "global_api_token_cap":self.config["global_api_token_cap"],"episode_api_token_cap":8_000_000, "context_tokens_are_estimated":True,"slurm_job_id":int(os.environ.get("SLURM_JOB_ID","0"))}) def emit(self, run_id, kind, *, metrics=None, status=None, config=None, message=None): run = next((r for r in self.config["runs"] if r["run_id"]==run_id),{}) event = dict(timestamp=datetime.now(timezone.utc).isoformat(),run_id=run_id, task=run.get("task"),condition=run.get("condition"),kind=kind) for key,value in [("metrics",metrics),("status",status),("config",config),("message",message)]: if value is not None: event[key]=value with self.event_lock: with self.events.open("a") as stream: stream.write(json.dumps(event)+"\n") def adapter(self, command, *options, output_log): if self.cancelled: raise InterruptedError("Study cancelled before adapter dispatch") argv = [self.args.evaluator_python,str(HERE/"programbench_adapter.py"),"--manifest",str(self.args.manifest.resolve()), "--programbench-root",str(self.args.programbench_root.resolve()),"--cpus","16","--memory","30g"] if self.args.resource_mode != "docker": argv += ["--resource-mode",self.args.resource_mode] if self.args.evaluator_wheelhouse is not None: argv += ["--evaluator-wheelhouse",str(self.args.evaluator_wheelhouse.resolve())] argv += [command,*map(str,options)] output_log.parent.mkdir(parents=True,exist_ok=True) with output_log.open("w") as stream: process = subprocess.Popen(argv,stdout=stream,stderr=subprocess.STDOUT,env=safe_env(),start_new_session=True) self.current=process self.current_is_adapter=True deadline=time.monotonic()+21600 terminated_at=None try: while process.poll() is None: if (self.cancelled or time.monotonic()>=deadline) and terminated_at is None: # SIGINT permits Python adapter context managers/finally # blocks to clean up evaluator resources before escalation. self.signal_process(process,signal.SIGINT,group=True) terminated_at=time.monotonic() elif terminated_at is not None and time.monotonic()-terminated_at>=5: self.signal_process(process,signal.SIGKILL,group=True) try: process.wait(timeout=0.25) except subprocess.TimeoutExpired: pass finally: if self.cancelled or terminated_at is not None or process.poll() is None: # Also kill descendants left behind by an exited adapter. self.signal_process(process,signal.SIGKILL,group=True) if process.poll() is None: process.wait(timeout=5) self.current=None self.current_is_adapter=False if self.cancelled: raise InterruptedError("Study cancelled during adapter execution") if terminated_at is not None: raise TimeoutError(f"Adapter {command} exceeded its time limit") lines = output_log.read_text(errors="replace").splitlines() result = None for line in reversed(lines): try: result=json.loads(line) break except json.JSONDecodeError: continue if process.returncode or not isinstance(result,dict): raise RuntimeError(f"Adapter {command} failed; inspect {output_log.name}") return result def prepare(self, run, instance_id, directory): return self.adapter("prepare","--instance-id",instance_id,"--episode-id",run["run_id"], "--output-dir",directory/"container",output_log=directory/"prepare.log") def runner(self, run, prepared, directory, prompt=None, wall_seconds=None): runner_dir=directory/"runner" config=dict(condition=run["condition"],run_id=run["run_id"],output_dir=str(runner_dir), proxy_url=f"http://127.0.0.1:{self.gateway.server_port}",task_prompt=prompt or prepared["task_prompt"], container_name=prepared["container_id"],container_user="agent",work_root=prepared["work_root"], seed_dir=prepared["seed_dir"],wall_seconds=wall_seconds or self.config["episode_wall_seconds"], snapshot_seconds=self.config["snapshot_seconds"],model_id=self.config["model"], max_output_tokens=16000,thinking_level="medium",temperature=0.7,top_p=0.9, agent_context_token_cap=run["agent_context_token_cap"],context_token_cap=run["context_token_cap"],max_active_agents=5) save(directory/"runner-config.json",config) with (directory/"runner.log").open("w") as log: self.current=subprocess.Popen(["node",str(HERE/"pi-runner/runner.mjs"),str(directory/"runner-config.json")], env=safe_env(self.local_key),stdout=log,stderr=subprocess.STDOUT,start_new_session=True) self.current_is_adapter=False deadline=time.monotonic()+config["wall_seconds"]+120 terminated_at=None while self.current.poll() is None: if (self.cancelled or time.monotonic()>=deadline) and terminated_at is None: self.current.send_signal(signal.SIGTERM) terminated_at=time.monotonic() elif terminated_at is not None and time.monotonic()-terminated_at>=30: self.signal_process(self.current,signal.SIGKILL,group=True) self.emit(run["run_id"],"progress",metrics={**self.gateway.ledger.totals(run["run_id"]), "study_elapsed_seconds":time.monotonic()-self.started}) try: self.current.wait(timeout=15) except subprocess.TimeoutExpired: pass code=self.current.returncode self.current=None result_path=runner_dir/"result.json" if not result_path.exists(): raise RuntimeError(f"Pi runner exited {code} without a result") return load(result_path) def cleanup_container(self, prepared): subprocess.run(["docker","rm","-f",prepared["container_id"]],capture_output=True,env=safe_env(),timeout=90) def settle_requests(self, run_id=None): deadline=time.monotonic()+660 while self.gateway.ledger.totals(run_id)["api_tokens_reserved"]: if time.monotonic()>deadline: raise RuntimeError("Provider usage settlement timed out; reservations remain charged") time.sleep(1) def validate(self): path=self.root/"validation.json" if path.exists(): prior=load(path) if prior.get("status")=="passed": return raise RuntimeError("Existing failed validation requires inspection; refusing automatic paid retry") reports=[] instance=self.manifest["tasks"][3]["instance_id"] for run in (r for r in self.config["runs"] if r["category"]=="validation"): directory=self.root/"validation"/run["run_id"] directory.mkdir(parents=True,exist_ok=True) self.emit(run["run_id"],"run_started",status="running",config={"model":self.config["model"],"kind":"infrastructure_validation"}) prepared=None try: prepared=self.prepare(run,instance,directory) prompt=("Infrastructure validation only. Ignore the reference executable. Implement a tiny independent command-line " "program: ./executable A B prints the integer sum of A and B and a newline. Provide executable compile.sh " "which creates executable without network access. Use Python standard library or shell. Commit and push " "to origin/main, test 2+3=5 and 10+(-4)=6, then finish promptly. ") if run["condition"]=="peers": prompt += "Coordinate briefly with at least one peer using send_message; divide implementation and checking without unnecessary changes." elif run["condition"]=="async": prompt += "Spawn a child to implement it, wait for the child's result, then resume/message that child to check the two examples. Integrate the working main branch and finish." result=self.runner(run,prepared,directory,prompt,300) script="set -eu; mkdir -p /workspace/validation-check; git --git-dir=/workspace/.pi-study/shared.git archive main | tar -x -C /workspace/validation-check; cd /workspace/validation-check; ./compile.sh; test \"$(./executable 2 3)\" = 5; test \"$(./executable 10 -4)\" = 6; test \"$(./executable 0 0)\" = 0" check=subprocess.run(["docker","exec","--user","agent",prepared["container_id"],"bash","-lc",script], capture_output=True,text=True,env=safe_env(),timeout=90) events=[json.loads(line) for line in (directory/"runner/events.jsonl").read_text().splitlines()] calls=sum(e["type"]=="model_request_end" for e in events) agent_count=len(result["agents"]) coordination=any(e["type"]=="message_sent" for e in events) idle_children=set() resume_messages=set() actual_resumptions=set() for event in events: if event["type"]=="agent_idle" and event.get("agent_id")!="agent-0": idle_children.add(event["agent_id"]) elif event["type"]=="message_sent" and event.get("agent_id")=="agent-0" and event.get("recipient_id") in idle_children: resume_messages.add(event["sequence"]) elif event["type"]=="message_delivered" and event.get("delivery")=="context_insertion": if event.get("causal_event_id") in resume_messages: actual_resumptions.add(event["agent_id"]) idle_children.discard(event.get("agent_id")) lifecycle_ok=(run["condition"]=="single" or coordination) and (run["condition"]!="peers" or agent_count==5) and (run["condition"]!="async" or (agent_count>=2 and bool(actual_resumptions))) passed=check.returncode==0 and calls>0 and lifecycle_ok and result["stop_reason"]=="completed" and not result.get("snapshot_error") report=dict(run_id=run["run_id"],passed=passed,model_requests=calls,agent_count=agent_count, coordinated=coordination,resumed_children=sorted(actual_resumptions),artifact_tests_passed=check.returncode==0,stop_reason=result["stop_reason"]) reports.append(report) self.emit(run["run_id"],"run_finished",status="passed" if passed else "failed", metrics={"validation_passed":int(passed),"artifact_tests_passed":int(check.returncode==0),"agents":agent_count}) if not passed: raise RuntimeError(f"Live Pi validation failed for {run['condition']}") except Exception as error: reports.append({"run_id":run["run_id"],"passed":False,"error":str(error)}) save(path,{"status":"failed","runs":reports}) raise finally: if prepared: self.cleanup_container(prepared) save(path,{"status":"passed","runs":reports}) def grade(self, run, directory, result): snapshots=sorted(result["snapshots"],key=lambda s:s["elapsed_ms"]) selected=[] # Grade saved scheduled states and the final artifact. Main-update # snapshots permit last-observed interpolation without hidden feedback. for cutoff in self.config["snapshot_seconds"]: eligible=[s for s in snapshots if s["elapsed_ms"]<=cutoff*1000] if eligible: selected.append((cutoff,eligible[-1])) if snapshots: selected.append((result["elapsed_seconds"],snapshots[-1])) grades=[] cache={} for cutoff,snapshot in sorted(selected,key=lambda item:item[0]): if self.cancelled: break archive=Path(snapshot["path"]) digest=hashlib.sha256(archive.read_bytes()).hexdigest() evaluation=directory/"evaluations"/digest summary_path=evaluation/run["task"]/"summary.json" if digest not in cache: try: summary=load(summary_path) if summary_path.exists() else self.adapter("grade","--instance-id",run["task"], "--submission",archive,"--calibration",self.args.calibration,"--output-dir",evaluation, output_log=evaluation/"grader.log") cache[digest]=summary except Exception as error: cache[digest]={"score":None,"complete":False,"error_code":"controller_grading_failure","error_details":str(error)} summary=cache[digest] # The adapter classifies candidate failures only after checking mask # coverage and infrastructure errors; do not infer zeros from a name. score=summary.get("analysis_score",summary.get("score")) valid=bool(summary.get("valid",summary.get("complete",False))) and isinstance(score,(int,float)) and not isinstance(score,bool) and math.isfinite(score) and 0<=score<=1 row={"snapshot_path":str(archive),"sha256":digest,"elapsed_seconds":float(cutoff), "snapshot_elapsed_seconds":snapshot["elapsed_ms"]/1000,"score":score if valid else None, "valid":valid,"evaluation_complete":bool(summary.get("complete")),"scoring_status":summary.get("scoring_status"),"error":summary.get("error_code"),"summary_path":str(summary_path)} grades.append(row) save(directory/"grades.json",grades) metrics={"elapsed_seconds":float(cutoff),"grader_complete":int(row["evaluation_complete"]),"grading_valid":int(row["valid"])} if row["valid"] and row["score"] is not None: metrics.update(hidden_test_fraction=float(row["score"]),tests_passed=summary.get("passed",0),tests_total=summary.get("test_count",0)) self.emit(run["run_id"],"score",metrics=metrics) return grades def run(self): if not self.args.validate_only: calibration=benchmark.check_calibration(self.args.calibration,self.args.manifest,self.manifest) if self.args.evaluator_wheelhouse is None: raise ValueError("--evaluator-wheelhouse is required before benchmark inference") dependency_digest=benchmark.verify_wheelhouse(self.args.evaluator_wheelhouse) if not dependency_digest or calibration.get("evaluator_dependencies_sha256")!=dependency_digest: raise ValueError("Evaluator dependency cache differs from calibration; refusing benchmark inference") reference_tasks=calibration.get("tasks",[]) if calibration.get("repetitions")!=2 or len(reference_tasks)!=len(self.manifest["tasks"]): raise ValueError("Two complete reference calibrations per selected task are required") for task in reference_tasks: repetitions=task.get("repetitions",[]) if len(repetitions)!=2: raise ValueError("Two complete reference calibrations per selected task are required") for reference in repetitions: score=reference.get("score") if reference.get("complete") is not True or not isinstance(score,(int,float)) or isinstance(score,bool) or not math.isfinite(score) or not 0.9<=score<=1: raise ValueError("Every reference calibration must complete with score at least 0.9") if reference.get("evaluator_dependencies_sha256")!=dependency_digest: raise ValueError("Reference repetition dependency cache differs from the pinned evaluator") self.emit("study-summary","progress",status="running",metrics={"reference_calibration_passed":1}, config={"evaluator_dependencies_sha256":dependency_digest}) self.validate() if self.args.validate_only: self.settle_requests() self.emit("study-summary","progress",status="validated",metrics={"live_pi_validation_passed":1}) return {"status":"interrupted" if self.cancelled else "validated"} for run in (r for r in self.config["runs"] if r["category"]=="benchmark"): if self.cancelled: break directory=self.root/"episodes"/run["run_id"] status_path=directory/"status.json" if status_path.exists(): prior=load(status_path) if prior.get("status") in {"completed","failed","interrupted","budget_exhausted"}: continue # A process crash must not silently grant a second attempt. save(status_path,{"status":"interrupted","stop_reason":"controller_restart"}) self.emit(run["run_id"],"run_finished",status="interrupted") continue directory.mkdir(parents=True,exist_ok=True) save(status_path,{"status":"running","run":run}) self.emit(run["run_id"],"run_started",status="running",config={**run,"model":self.config["model"],"reasoning":"medium"}) prepared=None try: prepared=self.prepare(run,run["task"],directory) result=self.runner(run,prepared,directory) self.cleanup_container(prepared) prepared=None self.settle_requests(run["run_id"]) grades=self.grade(run,directory,result) grading_failed=not grades or any(not g.get("valid") for g in grades) status="interrupted" if self.cancelled else "failed" if result["stop_reason"] in {"agent_error","setup_error"} or result.get("snapshot_error") or grading_failed else "completed" save(status_path,{"status":status,"stop_reason":result["stop_reason"],"graded_snapshots":len(grades),"grading_failed":grading_failed}) self.emit(run["run_id"],"run_finished",status=status,config={"stop_reason":result["stop_reason"]}, metrics={**self.gateway.ledger.totals(run["run_id"]),"elapsed_seconds":result["elapsed_seconds"]}) if result["stop_reason"] in {"global_api_budget","category_api_budget","daily_quota_headroom"}: break except Exception as error: status="interrupted" if self.cancelled else "failed" save(status_path,{"status":status,"error":str(error)}) self.emit(run["run_id"],"error",status=status,message=type(error).__name__) finally: if prepared: self.cleanup_container(prepared) self.settle_requests() if (HERE/"analyze.py").exists(): with (self.root/"analysis.log").open("w") as log: analysis=subprocess.run([sys.executable,str(HERE/"analyze.py"),"--experiment-dir",str(self.root),"--config",str(self.args.config.resolve())],stdout=log,stderr=subprocess.STDOUT,env=safe_env()) self.emit("study-summary","progress",metrics={"analysis_completed":int(analysis.returncode==0)}) planned=[r for r in self.config["runs"] if r["category"]=="benchmark"] statuses={r["run_id"]:load(self.root/"episodes"/r["run_id"]/"status.json") for r in planned if (self.root/"episodes"/r["run_id"]/"status.json").exists()} all_finished=len(statuses)==len(planned) failures=sum(s.get("status") in {"failed","interrupted"} for s in statuses.values()) overall="interrupted" if self.cancelled else "completed" if all_finished and not failures else "failed" if failures else "budget_exhausted" summary={"status":overall,"planned_episodes":len(planned),"recorded_episodes":len(statuses),"failed_episodes":failures, "usage":self.gateway.ledger.totals(),"runs":statuses} save(self.root/"study-status.json",summary) self.emit("study-summary","run_finished",status=overall,metrics={**self.gateway.ledger.totals(), "planned_episodes":len(planned),"recorded_episodes":len(statuses),"failed_episodes":failures}) return summary @staticmethod def signal_process(process, sig, *, group=False): try: if group: os.killpg(process.pid,sig) else: process.send_signal(sig) except ProcessLookupError: pass def stop(self,*_): self.cancelled=True if self.current: adapter=getattr(self,"current_is_adapter",False) self.signal_process(self.current,signal.SIGINT if adapter else signal.SIGTERM,group=adapter) def close(self): self.gateway.shutdown() self.gateway.server_close() # Wait for in-flight usage settlement. self.tracker.terminate() try: self.tracker.wait(timeout=60) except subprocess.TimeoutExpired: self.tracker.kill() self.tracker.wait() self.tracking_log.close() save(self.root/"usage-summary.json",self.gateway.ledger.totals()) def main(): parser=argparse.ArgumentParser(description=__doc__) parser.add_argument("--config",type=Path,default=HERE/"config.json") parser.add_argument("--manifest",type=Path,default=HERE.parent/"docs/programbench-five-task-manifest.json") parser.add_argument("--programbench-root",type=Path,required=True) parser.add_argument("--evaluator-python",default=sys.executable) parser.add_argument("--evaluator-wheelhouse",type=Path, default=Path(os.environ["PROGRAMBENCH_EVALUATOR_WHEELHOUSE"]) if os.environ.get("PROGRAMBENCH_EVALUATOR_WHEELHOUSE") else None) parser.add_argument("--calibration",type=Path,required=True) parser.add_argument("--output-dir",type=Path,required=True) parser.add_argument("--secrets-dir",type=Path,required=True) parser.add_argument("--resource-mode",choices=["docker","slurm"],default="docker") parser.add_argument("--validate-only",action="store_true") args=parser.parse_args() study=Study(args) signal.signal(signal.SIGTERM,study.stop) signal.signal(signal.SIGINT,study.stop) try: result=study.run() except BaseException as error: study.emit("study-summary","error",status="interrupted" if study.cancelled else "failed",message=type(error).__name__) if study.cancelled: result={"status":"interrupted"} else: raise finally: study.close() status=(result or {}).get("status") return 130 if status=="interrupted" else 1 if status=="failed" else 0 if __name__=="__main__": raise SystemExit(main())