#!/usr/bin/env python3 """Run source first, select the next intervention, then annotate the 2x2 surface.""" from __future__ import annotations import argparse import json import sys from pathlib import Path from typing import Any, Mapping HERE = Path(__file__).resolve().parent ACTION_DIR = HERE.parent / "action-aware-v0" sys.path.insert(0, str(ACTION_DIR)) sys.path.insert(0, str(HERE)) import pilot_controller as action_controller # noqa: E402 import prospective_decision # noqa: E402 SCHEMA = "active-intervention-prospective-state-v0" def validate_inputs(args: argparse.Namespace, manifest: Mapping[str, Any]) -> None: if manifest.get("schema") != "active-intervention-prospective-manifest-v0": raise RuntimeError("unexpected active intervention manifest schema") if manifest.get("status") != "PASS" or manifest["sanity"]["red_flags"]: raise RuntimeError("active intervention manifest did not pass preflight") required = { "manifest": args.manifest, "policy": args.policy, "aituner_root": args.aituner_root, "vllm_source": args.vllm_source, "venv_python": args.venv / "bin/python", "venv_vllm": args.venv / "bin/vllm", "model": args.model, "client": args.client, "burnin_study": Path(manifest["burnin"]["study"]), } for repetition, item in manifest["repetitions"].items(): required[f"rep{repetition}_study"] = Path(item["study"]) required[f"rep{repetition}_trace"] = Path(item["trace"]["path"]) missing = {name: str(path) for name, path in required.items() if not path.exists()} if missing: raise RuntimeError(f"active intervention input paths missing: {missing}") if prospective_decision.sha256_file(args.policy) != manifest["policy"]["sha256"]: raise RuntimeError("active intervention policy hash mismatch") def dry_run(args: argparse.Namespace, manifest: Mapping[str, Any]) -> dict[str, Any]: plan = action_controller.dry_run_plan(args, manifest) return { "schema": "active-intervention-prospective-dry-run-v0", "status": "PASS", "manifest": str(args.manifest), "policy": str(args.policy), "source_first": manifest["source_config_id"], "post_source_order": "selected by telemetry policy; all remaining cells then annotated", "candidate_actions": manifest["actions"], "projected_h20_hours": plan["projected_h20_hours"], "hard_cap_h20_hours": plan["hard_cap_h20_hours"], "sessions": plan["sessions"], } def load_or_build_decision( *, args: argparse.Namespace, run_root: Path ) -> dict[str, Any]: path = run_root / "active-decision.json" if path.exists(): decision = json.loads(path.read_text(encoding="utf-8")) if decision.get("manifest_sha256") != prospective_decision.sha256_file( args.manifest ): raise RuntimeError("existing active decision has a different manifest") if decision.get("policy_sha256") != prospective_decision.sha256_file(args.policy): raise RuntimeError("existing active decision has a different policy") return decision decision = prospective_decision.build_decision( manifest_path=args.manifest, policy_path=args.policy, run_root=run_root, ) prospective_decision.atomic_json(path, decision) return decision def parser() -> argparse.ArgumentParser: result = argparse.ArgumentParser() result.add_argument("--manifest", type=Path, required=True) result.add_argument("--policy", type=Path, required=True) result.add_argument("--run-root", type=Path, required=True) result.add_argument("--aituner-root", type=Path, required=True) result.add_argument("--vllm-source", type=Path, required=True) result.add_argument("--venv", type=Path, required=True) result.add_argument("--model", type=Path, required=True) result.add_argument("--client", type=Path, required=True) result.add_argument("--dry-run", action="store_true") return result def main() -> None: args = parser().parse_args() manifest = json.loads(args.manifest.read_text(encoding="utf-8")) validate_inputs(args, manifest) action_controller.configure(args, manifest) action_controller.base.MARKER = "active-intervention-prospective-v0" if args.dry_run: print(json.dumps(dry_run(args, manifest), indent=2, sort_keys=True)) return args.run_root.mkdir(parents=True, exist_ok=True) copied_manifest = args.run_root / "prospective-manifest.json" if not copied_manifest.exists(): action_controller.atomic_json(copied_manifest, manifest) state_path = args.run_root / "controller-state.json" state = action_controller.load_state( state_path, float(manifest["budget"]["hard_cap_h20_hours"]) ) state["schema"] = SCHEMA state["status"] = "running" action_controller.atomic_json(state_path, state) configs = {str(item["id"]): dict(item) for item in manifest["configs"]} config_indexes = { str(item["id"]): index for index, item in enumerate(manifest["configs"]) } source_id = str(manifest["source_config_id"]) action_controller.execute_session( args=args, manifest=manifest, config=configs[source_id], index=config_indexes[source_id], state=state, state_path=state_path, ) decision = load_or_build_decision(args=args, run_root=args.run_root) state["active_decision"] = { "path": str(args.run_root / "active-decision.json"), "status": decision["status"], "outcome_only": { key: decision["decisions"]["outcome_only"][key] for key in ("selected_cutoff_s", "decision_kind", "selected_action") }, "telemetry": { key: decision["decisions"]["telemetry"][key] for key in ("selected_cutoff_s", "decision_kind", "selected_action") }, } action_controller.atomic_json(state_path, state) if decision["status"] != "SELECTED": state["status"] = decision["status"].lower() state["completed_at"] = action_controller.time.time() action_controller.atomic_json(state_path, state) action_controller.wait_all_idle() print(json.dumps({"status": state["status"], "decision": decision["status"]})) return action_order = decision["decisions"]["telemetry"]["intervention_order"] execution_order = [source_id] for action_id in action_order: target_id = str(manifest["actions"][action_id]) if target_id not in execution_order: execution_order.append(target_id) for config_id in configs: if config_id not in execution_order: execution_order.append(config_id) state["execution_order"] = execution_order action_controller.atomic_json(state_path, state) for config_id in execution_order[1:]: action_controller.execute_session( args=args, manifest=manifest, config=configs[config_id], index=config_indexes[config_id], state=state, state_path=state_path, ) state["status"] = "complete" state["completed_at"] = action_controller.time.time() action_controller.atomic_json(state_path, state) action_controller.wait_all_idle() print( json.dumps( { "status": state["status"], "completed_sessions": state["completed_sessions"], "gpu_hours_total": state["gpu_hours_total"], "execution_order": execution_order, }, sort_keys=True, ) ) if __name__ == "__main__": main()