source/study/run_study.py

#!/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())