629 lines
21 KiB
Python
629 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}")
|
|
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=2.125 "
|
|
f"duration_s=300 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()
|