199 lines
7.6 KiB
Python
199 lines
7.6 KiB
Python
#!/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()
|