gahow/aituner
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add dry run gate for long telemetry pilot

Browse Source
This commit is contained in:
Gahow Wang
2026-07-14 17:25:45 +08:00
parent 52a9dc13dd
commit 2afc6eeb8d
2 changed files with 120 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

105
runs/intervention-response-v2/pilot_controller.py
View File

@@ -62,6 +62,100 @@ def configure(args: argparse.Namespace, manifest: dict[str, Any]) -> None:
}
def validate_inputs(args: argparse.Namespace, manifest: dict[str, Any]) -> None:
if manifest.get("schema") != "intervention-response-phase-aware-pilot-manifest-v2":
raise RuntimeError("unexpected phase-aware pilot manifest schema")
if manifest.get("status") != "PASS":
raise RuntimeError("phase-aware pilot manifest did not pass preflight")
failed_invariants = [
name
for name, passed in manifest.get("sanity", {}).get("invariants", {}).items()
if not passed
]
if failed_invariants:
raise RuntimeError(f"phase-aware pilot invariants failed: {failed_invariants}")
required = {
"manifest": args.manifest,
"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 replicate, repetition in manifest["repetitions"].items():
required[f"rep{replicate}_study"] = Path(repetition["study"])
required[f"rep{replicate}_trace"] = Path(
repetition["merged_trace"]["path"]
)
missing = {name: str(path) for name, path in required.items() if not path.exists()}
if missing:
raise RuntimeError(f"phase-aware pilot input paths missing: {missing}")
def dry_run_plan(args: argparse.Namespace, manifest: dict[str, Any]) -> dict[str, Any]:
sessions = []
for index, session in enumerate(manifest["sessions"]):
cell = f"tp4_mns{int(session['mns'])}"
entry = {"cell": cell, "gpus": (0, 1, 2, 3), "port": 8950 + index}
repetition = manifest["repetitions"][str(session["replicate"])]
session_root = args.run_root / "sessions" / str(session["session"])
high = repetition["selections"]["high"]
commands = {
"server": base.server_command(cell, entry["gpus"], entry["port"]),
"warmup": client_command(
entry,
study=repetition["study"],
anchor=float(high["anchor"]),
output=session_root / "warmup",
warmup=True,
),
"burnin": client_command(
entry,
study=manifest["burnin"]["study"],
anchor=float(manifest["burnin"]["anchor"]),
output=session_root / "burnin",
warmup=False,
),
}
for level in repetition["load_order"]:
selection = repetition["selections"][level]
commands[level] = client_command(
entry,
study=repetition["study"],
anchor=float(selection["anchor"]),
output=session_root / level,
warmup=False,
)
sessions.append(
{
"session": session["session"],
"replicate": int(session["replicate"]),
"mns": int(session["mns"]),
"port": entry["port"],
"load_order": repetition["load_order"],
"remaining_projection_h20_hours": remaining_projection(
len(manifest["sessions"]), index
),
"commands": {
role: shlex.join(command) for role, command in commands.items()
},
}
)
return {
"schema": "intervention-response-phase-aware-pilot-dry-run-v2",
"status": "PASS",
"manifest": str(args.manifest),
"run_root": str(args.run_root),
"session_count": len(sessions),
"projected_h20_hours": remaining_projection(len(sessions), 0),
"hard_cap_h20_hours": float(manifest["budget"]["hard_cap_h20_hours"]),
"sessions": sessions,
}
def load_state(path: Path, hard_cap: float) -> dict[str, Any]:
if path.exists():
return json.loads(path.read_text(encoding="utf-8"))
@@ -423,21 +517,22 @@ def parser() -> argparse.ArgumentParser:
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"))
if manifest.get("schema") != "intervention-response-phase-aware-pilot-manifest-v2":
raise RuntimeError("unexpected phase-aware pilot manifest schema")
if manifest["status"] != "PASS":
raise RuntimeError("phase-aware pilot manifest did not pass preflight")
validate_inputs(args, manifest)
configure(args, manifest)
if args.dry_run:
print(json.dumps(dry_run_plan(args, manifest), indent=2, sort_keys=True))
return
args.run_root.mkdir(parents=True, exist_ok=True)
copied_manifest = args.run_root / "pilot-manifest.json"
if not copied_manifest.exists():
atomic_json(copied_manifest, manifest)
configure(args, manifest)
state_path = args.run_root / "controller-state.json"
state = load_state(state_path, base.GPU_LIMIT)
state["status"] = "running"

20
runs/intervention-response-v2/test_analysis.py
View File

@@ -97,6 +97,26 @@ def main() -> None:
controller = load_controller_module()
assert math.isclose(controller.remaining_projection(6, 0), 7.7)
assert math.isclose(controller.remaining_projection(6, 5), 1.45)
parsed = controller.parser().parse_args(
[
"--manifest",
"/tmp/manifest.json",
"--run-root",
"/tmp/run",
"--aituner-root",
"/tmp/aituner",
"--vllm-source",
"/tmp/vllm",
"--venv",
"/tmp/venv",
"--model",
"/tmp/model",
"--client",
"/tmp/client.py",
"--dry-run",
]
)
assert parsed.dry_run is True
pilot_analysis = load_pilot_analysis_module()
stable = pilot_analysis.stable_adjacent_features(
[