From 2afc6eeb8dde91ae1280f78fef40c9f4cc33ba71 Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Tue, 14 Jul 2026 17:25:45 +0800 Subject: [PATCH] Add dry run gate for long telemetry pilot --- .../pilot_controller.py | 105 +++++++++++++++++- .../intervention-response-v2/test_analysis.py | 20 ++++ 2 files changed, 120 insertions(+), 5 deletions(-) diff --git a/runs/intervention-response-v2/pilot_controller.py b/runs/intervention-response-v2/pilot_controller.py index 98b5903..cb089ff 100644 --- a/runs/intervention-response-v2/pilot_controller.py +++ b/runs/intervention-response-v2/pilot_controller.py @@ -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" diff --git a/runs/intervention-response-v2/test_analysis.py b/runs/intervention-response-v2/test_analysis.py index ffd3a19..932dec4 100644 --- a/runs/intervention-response-v2/test_analysis.py +++ b/runs/intervention-response-v2/test_analysis.py @@ -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( [