Add prospective active intervention experiment

runs/active-intervention-v0/prospective_controller.py

@@ -0,0 +1,198 @@
+#!/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()