Make telemetry audit replay-phase aware
This commit is contained in:
471
runs/intervention-response-v2/pilot_controller.py
Normal file
471
runs/intervention-response-v2/pilot_controller.py
Normal file
@@ -0,0 +1,471 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Serialized, resumable controller for the 300-second phase-aware pilot."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import shlex
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
HERE = Path(__file__).resolve().parent
|
||||
PHASE6 = HERE.parent / "opprof-phase6"
|
||||
sys.path.insert(0, str(PHASE6))
|
||||
|
||||
import opprof_phase6_controller as base # noqa: E402
|
||||
|
||||
|
||||
SCHEMA = "intervention-response-phase-aware-pilot-state-v2"
|
||||
SESSION_ESTIMATE_H20_HOURS = 1.25
|
||||
SAFETY_H20_HOURS = 0.20
|
||||
CLIENT_TIMEOUT_S = 450.0
|
||||
|
||||
|
||||
def atomic_json(path: Path, payload: Any) -> None:
|
||||
base.atomic_json(path, payload)
|
||||
|
||||
|
||||
def wait_all_idle(timeout_s: float = 30.0) -> None:
|
||||
deadline = time.monotonic() + timeout_s
|
||||
last_error: Exception | None = None
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
base.assert_all_idle()
|
||||
return
|
||||
except RuntimeError as error:
|
||||
last_error = error
|
||||
time.sleep(1.0)
|
||||
raise last_error or RuntimeError("GPU idle timeout")
|
||||
|
||||
|
||||
def configure(args: argparse.Namespace, manifest: dict[str, Any]) -> None:
|
||||
base.WORKDIR = args.run_root.parent
|
||||
base.RUN_ROOT = args.run_root
|
||||
base.STATE = args.run_root / "controller-state.json"
|
||||
base.SOURCE = args.vllm_source
|
||||
base.VENV = args.venv
|
||||
base.AITUNER = args.aituner_root
|
||||
base.MODEL = args.model
|
||||
base.CLIENT = args.client
|
||||
base.GPU_LIMIT = float(manifest["budget"]["hard_cap_h20_hours"])
|
||||
base.MARKER = "intervention-response-phase-aware-v2"
|
||||
base.CELLS = {
|
||||
f"tp4_mns{mns}": {"tp": 4, "mns": int(mns)}
|
||||
for mns in manifest["engine"]["mns_endpoints"]
|
||||
}
|
||||
|
||||
|
||||
def load_state(path: Path, hard_cap: float) -> dict[str, Any]:
|
||||
if path.exists():
|
||||
return json.loads(path.read_text(encoding="utf-8"))
|
||||
return {
|
||||
"schema": SCHEMA,
|
||||
"status": "initialized",
|
||||
"hard_cap_h20_hours": hard_cap,
|
||||
"gpu_hours_total": 0.0,
|
||||
"completed_sessions": 0,
|
||||
"sessions": {},
|
||||
"failures": [],
|
||||
"started_at": time.time(),
|
||||
}
|
||||
|
||||
|
||||
def save_state(path: Path, state: dict[str, Any]) -> None:
|
||||
atomic_json(path, state)
|
||||
|
||||
|
||||
def append_echo(run_root: Path, line: str) -> None:
|
||||
run_root.mkdir(parents=True, exist_ok=True)
|
||||
with (run_root / "launch-echo.log").open("a", encoding="utf-8") as target:
|
||||
target.write(line + "\n")
|
||||
print(line, flush=True)
|
||||
|
||||
|
||||
def remaining_projection(session_count: int, index: int) -> float:
|
||||
return (session_count - index) * SESSION_ESTIMATE_H20_HOURS + SAFETY_H20_HOURS
|
||||
|
||||
|
||||
def start_server(
|
||||
*, session: dict[str, Any], index: int, run_root: Path
|
||||
) -> dict[str, Any]:
|
||||
cell = f"tp4_mns{int(session['mns'])}"
|
||||
gpus = (0, 1, 2, 3)
|
||||
session_root = run_root / "sessions" / str(session["session"])
|
||||
session_root.mkdir(parents=True, exist_ok=True)
|
||||
port = 8950 + index
|
||||
command = base.server_command(cell, gpus, port)
|
||||
with (session_root / "commands.log").open("a", encoding="utf-8") as log:
|
||||
log.write(f"SERVER {shlex.join(command)}\n")
|
||||
server_log = (session_root / "server.log").open("ab", buffering=0)
|
||||
environment = os.environ.copy()
|
||||
environment.update(
|
||||
{
|
||||
"CUDA_VISIBLE_DEVICES": "0,1,2,3",
|
||||
"VLLM_OPPROF_DIR": str(session_root / "opprof"),
|
||||
"OPPROF_PHASE6_MARKER": base.MARKER,
|
||||
"AITUNER_ROOT": str(base.AITUNER),
|
||||
"HF_HUB_OFFLINE": "1",
|
||||
"TRANSFORMERS_OFFLINE": "1",
|
||||
"PYTHONUNBUFFERED": "1",
|
||||
}
|
||||
)
|
||||
server = subprocess.Popen(
|
||||
command,
|
||||
cwd=base.SOURCE,
|
||||
env=environment,
|
||||
stdout=server_log,
|
||||
stderr=subprocess.STDOUT,
|
||||
start_new_session=True,
|
||||
)
|
||||
base.OWNED_PGIDS.add(server.pid)
|
||||
return {
|
||||
"cell": cell,
|
||||
"gpus": gpus,
|
||||
"port": port,
|
||||
"dir": session_root,
|
||||
"server": server,
|
||||
"server_handle": server_log,
|
||||
"spawned_at": time.time(),
|
||||
"results": [],
|
||||
}
|
||||
|
||||
|
||||
def client_command(
|
||||
entry: dict[str, Any],
|
||||
*,
|
||||
study: str,
|
||||
anchor: float,
|
||||
output: Path,
|
||||
warmup: bool,
|
||||
) -> list[str]:
|
||||
config = base.CELLS[entry["cell"]]
|
||||
command = [
|
||||
"taskset",
|
||||
"-c",
|
||||
base.cpu_mask(entry["gpus"]),
|
||||
str(base.VENV / "bin/python"),
|
||||
str(base.CLIENT),
|
||||
"warmup" if warmup else "run-anchor",
|
||||
"--study",
|
||||
study,
|
||||
"--cell",
|
||||
entry["cell"],
|
||||
"--anchor",
|
||||
str(anchor),
|
||||
"--tp",
|
||||
str(config["tp"]),
|
||||
"--mns",
|
||||
str(config["mns"]),
|
||||
"--base-url",
|
||||
f"http://127.0.0.1:{entry['port']}",
|
||||
"--result-dir",
|
||||
str(output),
|
||||
"--disable-slo-early-stop",
|
||||
]
|
||||
return command
|
||||
|
||||
|
||||
def run_client(
|
||||
*,
|
||||
entry: dict[str, Any],
|
||||
role: str,
|
||||
study: str,
|
||||
selection: dict[str, Any],
|
||||
output: Path,
|
||||
state: dict[str, Any],
|
||||
warmup: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
command = client_command(
|
||||
entry,
|
||||
study=study,
|
||||
anchor=float(selection["anchor"]),
|
||||
output=output,
|
||||
warmup=warmup,
|
||||
)
|
||||
with (entry["dir"] / "commands.log").open("a", encoding="utf-8") as log:
|
||||
log.write(f"CLIENT role={role} {shlex.join(command)}\n")
|
||||
handle = (output.parent / f"{output.name}.log").open("ab", buffering=0)
|
||||
environment = os.environ.copy()
|
||||
environment.update({"AITUNER_ROOT": str(base.AITUNER), "PYTHONUNBUFFERED": "1"})
|
||||
process = subprocess.Popen(
|
||||
command,
|
||||
cwd=base.WORKDIR,
|
||||
env=environment,
|
||||
stdout=handle,
|
||||
stderr=subprocess.STDOUT,
|
||||
start_new_session=True,
|
||||
)
|
||||
deadline = time.monotonic() + (180.0 if warmup else CLIENT_TIMEOUT_S)
|
||||
try:
|
||||
while process.poll() is None:
|
||||
if time.monotonic() > deadline:
|
||||
raise TimeoutError(f"client timeout: {entry['cell']} {role}")
|
||||
if entry["server"].poll() is not None:
|
||||
raise RuntimeError(f"server exited during {entry['cell']} {role}")
|
||||
base.assert_no_other_compute()
|
||||
if state["gpu_hours_total"] + base.live_gpu_hours([entry]) >= base.GPU_LIMIT:
|
||||
raise RuntimeError("phase-aware pilot H20-hour hard cap reached")
|
||||
time.sleep(1.0)
|
||||
except Exception:
|
||||
try:
|
||||
os.killpg(process.pid, signal.SIGTERM)
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
try:
|
||||
process.wait(timeout=10.0)
|
||||
except subprocess.TimeoutExpired:
|
||||
try:
|
||||
os.killpg(process.pid, signal.SIGKILL)
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
process.wait(timeout=10.0)
|
||||
raise
|
||||
finally:
|
||||
handle.close()
|
||||
if process.returncode:
|
||||
raise RuntimeError(
|
||||
f"client failed: cell={entry['cell']} role={role} rc={process.returncode}"
|
||||
)
|
||||
result = json.loads((output / "result.json").read_text(encoding="utf-8"))
|
||||
validate_result(
|
||||
result=result,
|
||||
selection=selection,
|
||||
role=role,
|
||||
warmup=warmup,
|
||||
)
|
||||
entry["results"].append(
|
||||
{
|
||||
"anchor": float(selection["anchor"]),
|
||||
"dir": str(output),
|
||||
"kind": result["kind"],
|
||||
}
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def validate_result(
|
||||
*, result: dict[str, Any], selection: dict[str, Any], role: str, warmup: bool
|
||||
) -> None:
|
||||
if result.get("slo_early_stop_disabled") is not True:
|
||||
raise RuntimeError(f"SLO early stop was not disabled: {role}")
|
||||
if warmup:
|
||||
if result["kind"] != "warmup" or int(result["selection"]["count"]) != 16:
|
||||
raise RuntimeError(f"invalid warmup result: {role}")
|
||||
if not all(
|
||||
result["invariants"].get(key, False)
|
||||
for key in ("warmup_16", "warmup_exact_16", "warmup_long")
|
||||
):
|
||||
raise RuntimeError(f"warmup invariant failed: {role}")
|
||||
return
|
||||
if bool(result["early_stopped"]):
|
||||
raise RuntimeError(f"uncensored run early-stopped: {role}")
|
||||
if int(result["selection"]["count"]) != int(selection["selected_count"]):
|
||||
raise RuntimeError(f"selection count mismatch: {role}")
|
||||
for key, manifest_key in (
|
||||
("request_id_order_sha256", "request_id_order_sha256"),
|
||||
("arrival_order_sha256", "arrival_order_sha256"),
|
||||
("raw_length_order_sha256", "input_length_order_sha256"),
|
||||
):
|
||||
if result["selection"][key] != selection[manifest_key]:
|
||||
raise RuntimeError(f"selection hash mismatch {key}: {role}")
|
||||
if int(result["observed_count"]) != int(selection["selected_count"]):
|
||||
raise RuntimeError(f"request accounting mismatch: {role}")
|
||||
|
||||
|
||||
def execute_session(
|
||||
*,
|
||||
index: int,
|
||||
session: dict[str, Any],
|
||||
manifest: dict[str, Any],
|
||||
run_root: Path,
|
||||
state_path: Path,
|
||||
state: dict[str, Any],
|
||||
) -> None:
|
||||
name = str(session["session"])
|
||||
if state["sessions"].get(name, {}).get("status") == "complete":
|
||||
return
|
||||
projection = remaining_projection(len(manifest["sessions"]), index)
|
||||
if state["gpu_hours_total"] + projection > base.GPU_LIMIT:
|
||||
state["status"] = "budget_projection_stop"
|
||||
state["budget_stop"] = {
|
||||
"before_session": name,
|
||||
"spent_h20_hours": state["gpu_hours_total"],
|
||||
"remaining_projection_h20_hours": projection,
|
||||
"hard_cap_h20_hours": base.GPU_LIMIT,
|
||||
}
|
||||
save_state(state_path, state)
|
||||
raise RuntimeError(f"projected pilot cost exceeds hard cap before {name}")
|
||||
|
||||
replicate = str(session["replicate"])
|
||||
repetition = manifest["repetitions"][replicate]
|
||||
echo = (
|
||||
f"PHASE_PILOT_SESSION_ECHO session={name} tp=4 mns={session['mns']} "
|
||||
f"gpus=0-3 workload={manifest['source']['window_id']} duration_s=300 "
|
||||
f"loads={','.join(repetition['load_order'])} disable_slo_early_stop=true "
|
||||
f"spent_h20h={state['gpu_hours_total']:.6f} "
|
||||
f"remaining_projection_h20h={projection:.3f} cap_h20h={base.GPU_LIMIT:.1f} "
|
||||
f"manifest={run_root / 'pilot-manifest.json'}"
|
||||
)
|
||||
append_echo(run_root, echo)
|
||||
wait_all_idle()
|
||||
session_state = {
|
||||
"status": "starting",
|
||||
"replicate": int(replicate),
|
||||
"mns": int(session["mns"]),
|
||||
"started_at": time.time(),
|
||||
"runs": [],
|
||||
}
|
||||
state["status"] = "running"
|
||||
state["sessions"][name] = session_state
|
||||
save_state(state_path, state)
|
||||
entry = start_server(session=session, index=index, run_root=run_root)
|
||||
failure: Exception | None = None
|
||||
try:
|
||||
base.wait_ready(entry)
|
||||
high = repetition["selections"]["high"]
|
||||
session_state["status"] = "warmup"
|
||||
save_state(state_path, state)
|
||||
run_client(
|
||||
entry=entry,
|
||||
role="warmup",
|
||||
study=repetition["study"],
|
||||
selection=high,
|
||||
output=entry["dir"] / "warmup",
|
||||
state=state,
|
||||
warmup=True,
|
||||
)
|
||||
session_state["status"] = "burnin"
|
||||
save_state(state_path, state)
|
||||
burnin = manifest["burnin"]
|
||||
burnin_result = run_client(
|
||||
entry=entry,
|
||||
role="burnin",
|
||||
study=burnin["study"],
|
||||
selection=burnin,
|
||||
output=entry["dir"] / "burnin",
|
||||
state=state,
|
||||
)
|
||||
session_state["burnin"] = {
|
||||
"pass_rate": burnin_result["pass_rate"],
|
||||
"feasible": burnin_result["feasible"],
|
||||
"elapsed_s": burnin_result["interval"]["elapsed_s"],
|
||||
}
|
||||
session_state["status"] = "measured"
|
||||
save_state(state_path, state)
|
||||
for level in repetition["load_order"]:
|
||||
selection = repetition["selections"][level]
|
||||
result = run_client(
|
||||
entry=entry,
|
||||
role=level,
|
||||
study=repetition["study"],
|
||||
selection=selection,
|
||||
output=entry["dir"] / level,
|
||||
state=state,
|
||||
)
|
||||
session_state["runs"].append(
|
||||
{
|
||||
"level": level,
|
||||
"selected_count": selection["selected_count"],
|
||||
"offered_req_s_per_gpu": selection["offered_req_s_per_gpu"],
|
||||
"pass_rate": result["pass_rate"],
|
||||
"feasible": result["feasible"],
|
||||
"elapsed_s": result["interval"]["elapsed_s"],
|
||||
"early_stopped": result["early_stopped"],
|
||||
}
|
||||
)
|
||||
save_state(state_path, state)
|
||||
session_state["status"] = "stopping"
|
||||
save_state(state_path, state)
|
||||
except Exception as error: # noqa: BLE001
|
||||
failure = error
|
||||
finally:
|
||||
try:
|
||||
base.stop_entry(entry)
|
||||
except Exception as error: # noqa: BLE001
|
||||
failure = failure or error
|
||||
time.sleep(2.0)
|
||||
try:
|
||||
wait_all_idle()
|
||||
except Exception as error: # noqa: BLE001
|
||||
failure = failure or error
|
||||
|
||||
session_hours = base.live_gpu_hours([entry])
|
||||
state["gpu_hours_total"] += session_hours
|
||||
session_state["gpu_hours"] = session_hours
|
||||
if failure is not None:
|
||||
session_state["status"] = "failed"
|
||||
session_state["failure"] = repr(failure)
|
||||
state["status"] = "failed"
|
||||
state["failures"].append({"session": name, "failure": repr(failure)})
|
||||
save_state(state_path, state)
|
||||
raise failure
|
||||
validation = base.validate_cell(entry)
|
||||
session_state["validation"] = validation
|
||||
session_state["status"] = "complete"
|
||||
session_state["completed_at"] = time.time()
|
||||
state["completed_sessions"] += 1
|
||||
save_state(state_path, state)
|
||||
|
||||
|
||||
def parser() -> argparse.ArgumentParser:
|
||||
result = argparse.ArgumentParser()
|
||||
result.add_argument("--manifest", 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)
|
||||
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")
|
||||
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"
|
||||
save_state(state_path, state)
|
||||
for index, session in enumerate(manifest["sessions"]):
|
||||
execute_session(
|
||||
index=index,
|
||||
session=session,
|
||||
manifest=manifest,
|
||||
run_root=args.run_root,
|
||||
state_path=state_path,
|
||||
state=state,
|
||||
)
|
||||
state["status"] = "complete"
|
||||
state["completed_at"] = time.time()
|
||||
save_state(state_path, state)
|
||||
wait_all_idle()
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"status": state["status"],
|
||||
"completed_sessions": state["completed_sessions"],
|
||||
"gpu_hours_total": state["gpu_hours_total"],
|
||||
},
|
||||
sort_keys=True,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user