Files
aituner/runs/action-aware-v0/pilot_controller.py

639 lines
21 KiB
Python

#!/usr/bin/env python3
"""Serialized controller for the crossed-constraint action-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, Mapping
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 = "action-aware-constraint-pilot-state-v0"
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: Mapping[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 = "action-aware-constraint-pilot-v0"
def validate_inputs(args: argparse.Namespace, manifest: Mapping[str, Any]) -> None:
if manifest.get("schema") not in {
"action-aware-constraint-pilot-manifest-v0",
"action-aware-constraint-pilot-manifest-v1",
}:
raise RuntimeError("unexpected action-aware manifest schema")
if manifest.get("status") != "PASS":
raise RuntimeError("action-aware manifest did not pass preflight")
red_flags = manifest.get("sanity", {}).get("red_flags", [])
if red_flags:
raise RuntimeError(f"manifest red flags: {red_flags}")
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 repetition, item in manifest["repetitions"].items():
required[f"rep{repetition}_study"] = Path(item["study"])
required[f"rep{repetition}_trace"] = Path(item["merged_trace"]["path"])
missing = {name: str(path) for name, path in required.items() if not path.exists()}
if missing:
raise RuntimeError(f"action-aware input paths missing: {missing}")
def config_map(manifest: Mapping[str, Any]) -> dict[str, dict[str, Any]]:
return {str(item["id"]): dict(item) for item in manifest["configs"]}
def server_command(
config: Mapping[str, Any], *, gpus: tuple[int, ...], port: int
) -> list[str]:
return [
"taskset",
"-c",
base.cpu_mask(gpus),
str(base.VENV / "bin/vllm"),
"serve",
str(base.MODEL),
"--host",
"127.0.0.1",
"--port",
str(port),
"--served-model-name",
"qwen3-30b-a3b-community",
"--max-num-batched-tokens",
str(config["mbbt"]),
"--max-num-seqs",
str(config["mns"]),
"--tensor-parallel-size",
"4",
"--shutdown-timeout",
"120",
]
def client_command(
entry: Mapping[str, Any],
config: Mapping[str, Any],
*,
study: str,
anchor: float,
output: Path,
warmup: bool,
) -> list[str]:
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",
str(config["id"]),
"--anchor",
str(anchor),
"--tp",
"4",
"--mns",
str(config["mns"]),
"--mbbt",
str(config["mbbt"]),
"--base-url",
f"http://127.0.0.1:{entry['port']}",
"--result-dir",
str(output),
"--disable-slo-early-stop",
]
return command
def remaining_projection(
manifest: Mapping[str, Any], *, completed_sessions: int
) -> float:
remaining = len(manifest["configs"]) - completed_sessions
return (
remaining * float(manifest["budget"]["session_estimate_h20_hours"])
+ float(manifest["budget"]["safety_h20_hours"])
)
def dry_run_plan(
args: argparse.Namespace, manifest: Mapping[str, Any]
) -> dict[str, Any]:
sessions = []
for index, config in enumerate(manifest["configs"]):
entry = {"gpus": (0, 1, 2, 3), "port": 9050 + index}
session_root = args.run_root / "sessions" / str(config["id"])
first_repetition = str(config["repetition_order"][0])
first = manifest["repetitions"][first_repetition]
commands = {
"server": server_command(config, gpus=entry["gpus"], port=entry["port"]),
"warmup": client_command(
entry,
config,
study=first["study"],
anchor=float(first["selection"]["anchor"]),
output=session_root / "warmup",
warmup=True,
),
"burnin": client_command(
entry,
config,
study=manifest["burnin"]["study"],
anchor=float(manifest["burnin"]["anchor"]),
output=session_root / "burnin",
warmup=False,
),
}
for repetition in config["repetition_order"]:
item = manifest["repetitions"][str(repetition)]
commands[f"rep{repetition}"] = client_command(
entry,
config,
study=item["study"],
anchor=float(item["selection"]["anchor"]),
output=session_root / f"rep{repetition}",
warmup=False,
)
sessions.append(
{
"config": config["id"],
"mns": config["mns"],
"mbbt": config["mbbt"],
"port": entry["port"],
"repetition_order": config["repetition_order"],
"commands": {
role: shlex.join(command) for role, command in commands.items()
},
}
)
return {
"schema": "action-aware-constraint-pilot-dry-run-v0",
"status": "PASS",
"manifest": str(args.manifest),
"run_root": str(args.run_root),
"projected_h20_hours": remaining_projection(
manifest, completed_sessions=0
),
"hard_cap_h20_hours": 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"))
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 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 start_server(
*,
args: argparse.Namespace,
config: Mapping[str, Any],
index: int,
) -> dict[str, Any]:
gpus = (0, 1, 2, 3)
session_root = args.run_root / "sessions" / str(config["id"])
session_root.mkdir(parents=True, exist_ok=True)
port = 9050 + index
command = server_command(config, gpus=gpus, port=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": str(config["id"]),
"gpus": gpus,
"port": port,
"dir": session_root,
"server": server,
"server_handle": server_log,
"spawned_at": time.time(),
"results": [],
}
def validate_result(
result: Mapping[str, Any],
*,
config: Mapping[str, Any],
selection: Mapping[str, Any],
role: str,
warmup: bool,
) -> None:
if result.get("schema") != "action-aware-pilot-result-v0":
raise RuntimeError(f"unexpected result schema: {role}")
if result.get("config_id") != config["id"]:
raise RuntimeError(f"config id mismatch: {role}")
if int(result["tp"]) != 4:
raise RuntimeError(f"TP mismatch: {role}")
if int(result["mns"]) != int(config["mns"]):
raise RuntimeError(f"MNS mismatch: {role}")
if int(result["mbbt"]) != int(config["mbbt"]):
raise RuntimeError(f"MBBT mismatch: {role}")
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: {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}")
if int(result["observed_count"]) != int(selection["selected_count"]):
raise RuntimeError(f"request accounting mismatch: {role}")
for result_key, selection_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"][result_key] != selection[selection_key]:
raise RuntimeError(f"selection hash mismatch {result_key}: {role}")
def burnin_gate(
result: Mapping[str, Any],
*,
expected_count: int,
maximum_elapsed_s: float,
) -> dict[str, Any]:
if result.get("kind") != "anchor":
raise RuntimeError("burnin gate received a non-anchor result")
if int(result["selection"]["count"]) != expected_count:
raise RuntimeError("burnin gate received the wrong request set")
elapsed_s = float(result["interval"]["elapsed_s"])
summary = {
"elapsed_s": elapsed_s,
"pass_rate": float(result["pass_rate"]),
"feasible": bool(result["feasible"]),
}
if elapsed_s > maximum_elapsed_s:
raise RuntimeError(
f"burnin throughput gate failed: {elapsed_s:.3f}s > "
f"{maximum_elapsed_s:.3f}s"
)
return summary
def run_client(
*,
entry: dict[str, Any],
config: Mapping[str, Any],
role: str,
study: str,
selection: Mapping[str, Any],
output: Path,
state: Mapping[str, Any],
timeout_s: float,
warmup: bool = False,
) -> dict[str, Any]:
command = client_command(
entry,
config,
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() + timeout_s
try:
while process.poll() is None:
if time.monotonic() > deadline:
raise TimeoutError(f"client timeout: {config['id']} {role}")
if entry["server"].poll() is not None:
raise RuntimeError(f"server exited during {config['id']} {role}")
base.assert_no_other_compute()
if state["gpu_hours_total"] + base.live_gpu_hours([entry]) >= base.GPU_LIMIT:
raise RuntimeError("action-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: config={config['id']} role={role} rc={process.returncode}"
)
result = json.loads((output / "result.json").read_text(encoding="utf-8"))
validate_result(
result,
config=config,
selection=selection,
role=role,
warmup=warmup,
)
entry["results"].append(
{"anchor": float(selection["anchor"]), "dir": str(output), "kind": result["kind"]}
)
return result
def execute_session(
*,
args: argparse.Namespace,
manifest: Mapping[str, Any],
config: Mapping[str, Any],
index: int,
state: dict[str, Any],
state_path: Path,
) -> None:
name = str(config["id"])
if state["sessions"].get(name, {}).get("status") == "complete":
return
projection = remaining_projection(
manifest, completed_sessions=int(state["completed_sessions"])
)
if float(state["gpu_hours_total"]) + projection > base.GPU_LIMIT:
raise RuntimeError(f"projected cost exceeds cap before {name}")
load_values = {
float(item["selection"]["offered_req_s_per_gpu"])
for item in manifest["repetitions"].values()
}
load_text = (
f"{next(iter(load_values)):.6g}"
if len(load_values) == 1
else ",".join(f"{value:.6g}" for value in sorted(load_values))
)
echo = (
f"ACTION_AWARE_SESSION_ECHO host=dash0 config={name} tp=4 "
f"mns={config['mns']} mbbt={config['mbbt']} gpus=0-3 "
f"workload={manifest['source']['window_id']} load_per_gpu={load_text} "
f"duration_s={manifest['engine']['duration_s']} "
f"repetitions={','.join(map(str, config['repetition_order']))} "
f"source={args.manifest} output={args.run_root / 'sessions' / name} "
f"spent_h20h={state['gpu_hours_total']:.6f} "
f"remaining_projection_h20h={projection:.3f} cap_h20h={base.GPU_LIMIT:.1f}"
)
append_echo(args.run_root, echo)
wait_all_idle()
session_state = {
"status": "starting",
"mns": int(config["mns"]),
"mbbt": int(config["mbbt"]),
"repetition_order": list(config["repetition_order"]),
"started_at": time.time(),
"runs": [],
}
state["status"] = "running"
state["sessions"][name] = session_state
atomic_json(state_path, state)
entry = start_server(args=args, config=config, index=index)
failure: Exception | None = None
try:
base.wait_ready(entry)
first = manifest["repetitions"][str(config["repetition_order"][0])]
session_state["status"] = "warmup"
atomic_json(state_path, state)
run_client(
entry=entry,
config=config,
role="warmup",
study=first["study"],
selection=first["selection"],
output=entry["dir"] / "warmup",
state=state,
timeout_s=180.0,
warmup=True,
)
session_state["status"] = "burnin"
atomic_json(state_path, state)
burnin = manifest["burnin"]
burnin_result = run_client(
entry=entry,
config=config,
role="burnin",
study=burnin["study"],
selection=burnin,
output=entry["dir"] / "burnin",
state=state,
timeout_s=float(manifest["engine"]["client_timeout_s"]),
)
session_state["burnin"] = burnin_gate(
burnin_result,
expected_count=int(burnin["selected_count"]),
maximum_elapsed_s=float(manifest["engine"]["burnin_max_elapsed_s"]),
)
atomic_json(state_path, state)
session_state["status"] = "measured"
atomic_json(state_path, state)
for repetition in config["repetition_order"]:
item = manifest["repetitions"][str(repetition)]
role = f"rep{repetition}"
result = run_client(
entry=entry,
config=config,
role=role,
study=item["study"],
selection=item["selection"],
output=entry["dir"] / role,
state=state,
timeout_s=float(manifest["engine"]["client_timeout_s"]),
)
session_state["runs"].append(
{
"repetition": int(repetition),
"pass_rate": result["pass_rate"],
"feasible": result["feasible"],
"slo_pass_count": result["slo_pass_count"],
"elapsed_s": result["interval"]["elapsed_s"],
}
)
atomic_json(state_path, state)
session_state["status"] = "stopping"
atomic_json(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)})
atomic_json(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
atomic_json(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)
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)
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)
state_path = args.run_root / "controller-state.json"
state = load_state(state_path, base.GPU_LIMIT)
state["status"] = "running"
atomic_json(state_path, state)
for index, config in enumerate(manifest["configs"]):
execute_session(
args=args,
manifest=manifest,
config=config,
index=index,
state=state,
state_path=state_path,
)
state["status"] = "complete"
state["completed_at"] = time.time()
atomic_json(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()