437 lines
14 KiB
Python
437 lines
14 KiB
Python
#!/usr/bin/env python3
|
|
"""Serialized dash0 controller for the exact-timestamp prefix pilot."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import shlex
|
|
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
|
|
|
|
|
|
ORDER = (
|
|
"tp1_mns8",
|
|
"tp1_mns64",
|
|
"tp2_mns8",
|
|
"tp2_mns64",
|
|
"tp4_mns16",
|
|
"tp4_mns64",
|
|
)
|
|
CELL_ESTIMATE_H20_HOURS = {1: 0.20, 2: 0.40, 4: 0.80}
|
|
SAFETY_H20_HOURS = 0.20
|
|
|
|
|
|
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_base(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["execution"]["hard_cap_h20_hours"])
|
|
base.MARKER = "fidelity-prefix-pilot-20260714"
|
|
base.CELLS = {
|
|
cell: {"tp": int(config["tp"]), "mns": int(config["mns"])}
|
|
for cell, config in manifest["cells"].items()
|
|
}
|
|
|
|
|
|
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": "fidelity-prefix-pilot-state-v1",
|
|
"status": "initialized",
|
|
"hard_cap_h20_hours": hard_cap,
|
|
"gpu_hours_total": 0.0,
|
|
"completed_cells": 0,
|
|
"cells": {},
|
|
"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(manifest: dict[str, Any], index: int) -> float:
|
|
return sum(
|
|
CELL_ESTIMATE_H20_HOURS[int(manifest["cells"][cell]["tp"])]
|
|
for cell in ORDER[index:]
|
|
) + SAFETY_H20_HOURS
|
|
|
|
|
|
def start_server(
|
|
*,
|
|
cell: str,
|
|
index: int,
|
|
run_root: Path,
|
|
) -> dict[str, Any]:
|
|
config = base.CELLS[cell]
|
|
gpus = tuple(range(int(config["tp"])))
|
|
cell_root = run_root / "cells" / cell
|
|
cell_root.mkdir(parents=True, exist_ok=True)
|
|
port = 8900 + index
|
|
command = base.server_command(cell, gpus, port)
|
|
with (cell_root / "commands.log").open("a", encoding="utf-8") as log:
|
|
log.write(f"SERVER {shlex.join(command)}\n")
|
|
server_log = (cell_root / "server.log").open("ab", buffering=0)
|
|
environment = os.environ.copy()
|
|
environment.update(
|
|
{
|
|
"CUDA_VISIBLE_DEVICES": ",".join(map(str, gpus)),
|
|
"VLLM_OPPROF_DIR": str(cell_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": cell_root,
|
|
"server": server,
|
|
"server_handle": server_log,
|
|
"spawned_at": time.time(),
|
|
"results": [],
|
|
}
|
|
|
|
|
|
def selection_for(
|
|
manifest: dict[str, Any], cell: str, role: str
|
|
) -> tuple[str, dict[str, Any]]:
|
|
level = "low" if role == "burnin" or role.startswith("low") else "high"
|
|
return level, manifest["cells"][cell]["targets"][level]["selections"][role]
|
|
|
|
|
|
def client_command(
|
|
entry: dict[str, Any],
|
|
*,
|
|
role: str,
|
|
selection: dict[str, Any],
|
|
output: Path,
|
|
warmup: bool,
|
|
) -> list[str]:
|
|
config = base.CELLS[entry["cell"]]
|
|
return [
|
|
"taskset",
|
|
"-c",
|
|
base.cpu_mask(entry["gpus"]),
|
|
str(base.VENV / "bin/python"),
|
|
str(base.CLIENT),
|
|
"warmup" if warmup else "run-anchor",
|
|
"--study",
|
|
str(selection["study"]),
|
|
"--cell",
|
|
entry["cell"],
|
|
"--anchor",
|
|
str(selection["anchor"]),
|
|
"--tp",
|
|
str(config["tp"]),
|
|
"--mns",
|
|
str(config["mns"]),
|
|
"--base-url",
|
|
f"http://127.0.0.1:{entry['port']}",
|
|
"--result-dir",
|
|
str(output),
|
|
]
|
|
|
|
|
|
def run_client(
|
|
*,
|
|
entry: dict[str, Any],
|
|
role: str,
|
|
selection: dict[str, Any],
|
|
output: Path,
|
|
state: dict[str, Any],
|
|
warmup: bool = False,
|
|
) -> dict[str, Any]:
|
|
command = client_command(
|
|
entry, role=role, selection=selection, 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
|
|
try:
|
|
while process.poll() is None:
|
|
if time.monotonic() > deadline:
|
|
process.terminate()
|
|
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:
|
|
process.terminate()
|
|
raise RuntimeError("pilot H20-hour hard cap reached")
|
|
time.sleep(1.0)
|
|
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"))
|
|
if int(result["selection"]["count"]) != int(selection["selected_count"]):
|
|
raise RuntimeError(f"selection count mismatch: {entry['cell']} {role}")
|
|
for key in (
|
|
"request_id_order_sha256",
|
|
"arrival_order_sha256",
|
|
"raw_length_order_sha256",
|
|
):
|
|
manifest_key = (
|
|
"input_length_order_sha256" if key == "raw_length_order_sha256" else key
|
|
)
|
|
if result["selection"][key] != selection[manifest_key]:
|
|
raise RuntimeError(f"selection hash mismatch {key}: {entry['cell']} {role}")
|
|
entry["results"].append(
|
|
{"anchor": float(selection["anchor"]), "dir": str(output), "kind": result["kind"]}
|
|
)
|
|
return result
|
|
|
|
|
|
def execute_cell(
|
|
*,
|
|
index: int,
|
|
cell: str,
|
|
manifest: dict[str, Any],
|
|
run_root: Path,
|
|
state_path: Path,
|
|
state: dict[str, Any],
|
|
) -> None:
|
|
if state["cells"].get(cell, {}).get("status") == "complete":
|
|
return
|
|
projection = remaining_projection(manifest, index)
|
|
if state["gpu_hours_total"] + projection > base.GPU_LIMIT:
|
|
state["status"] = "budget_projection_stop"
|
|
state["budget_stop"] = {
|
|
"before_cell": cell,
|
|
"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 {cell}")
|
|
|
|
config = manifest["cells"][cell]
|
|
echo = (
|
|
f"PILOT_CELL_ECHO cell={cell} tp={config['tp']} mns={config['mns']} "
|
|
f"gpus=0-{int(config['tp']) - 1} workload={manifest['source']['window_id']} "
|
|
f"roles=burnin+low1/high1/low2/high2/low3/high3 "
|
|
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()
|
|
cell_state = {
|
|
"status": "starting",
|
|
"tp": int(config["tp"]),
|
|
"mns": int(config["mns"]),
|
|
"started_at": time.time(),
|
|
"runs": [],
|
|
}
|
|
state["status"] = "running"
|
|
state["cells"][cell] = cell_state
|
|
save_state(state_path, state)
|
|
entry = start_server(cell=cell, index=index, run_root=run_root)
|
|
failure: Exception | None = None
|
|
try:
|
|
base.wait_ready(entry)
|
|
_level, burnin = selection_for(manifest, cell, "burnin")
|
|
cell_state["status"] = "warmup"
|
|
save_state(state_path, state)
|
|
warmup = run_client(
|
|
entry=entry,
|
|
role="burnin",
|
|
selection=burnin,
|
|
output=entry["dir"] / "warmup",
|
|
state=state,
|
|
warmup=True,
|
|
)
|
|
cell_state["warmup"] = {
|
|
"exact_output_count": warmup["exact_output_count"],
|
|
"long_gt4096": warmup["selection"]["long_gt4096"],
|
|
}
|
|
cell_state["status"] = "burnin"
|
|
save_state(state_path, state)
|
|
burnin_result = run_client(
|
|
entry=entry,
|
|
role="burnin",
|
|
selection=burnin,
|
|
output=entry["dir"] / "burnin",
|
|
state=state,
|
|
)
|
|
cell_state["burnin"] = {
|
|
"pass_rate": burnin_result["pass_rate"],
|
|
"feasible": burnin_result["feasible"],
|
|
}
|
|
role_order = manifest["execution"][
|
|
"even_cell_order" if index % 2 == 0 else "odd_cell_order"
|
|
]
|
|
cell_state["status"] = "measured"
|
|
cell_state["role_order"] = role_order
|
|
save_state(state_path, state)
|
|
for role in role_order:
|
|
level, selection = selection_for(manifest, cell, role)
|
|
result = run_client(
|
|
entry=entry,
|
|
role=role,
|
|
selection=selection,
|
|
output=entry["dir"] / f"{level}-rep{role[-1]}",
|
|
state=state,
|
|
)
|
|
cell_state["runs"].append(
|
|
{
|
|
"role": role,
|
|
"level": level,
|
|
"anchor": selection["anchor"],
|
|
"selected_count": selection["selected_count"],
|
|
"pass_rate": result["pass_rate"],
|
|
"feasible": result["feasible"],
|
|
"elapsed_s": result["interval"]["elapsed_s"],
|
|
}
|
|
)
|
|
save_state(state_path, state)
|
|
cell_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
|
|
|
|
cell_hours = base.live_gpu_hours([entry])
|
|
state["gpu_hours_total"] += cell_hours
|
|
cell_state["gpu_hours"] = cell_hours
|
|
if failure is not None:
|
|
cell_state["status"] = "failed"
|
|
cell_state["failure"] = repr(failure)
|
|
state["status"] = "failed"
|
|
state["failures"].append({"cell": cell, "failure": repr(failure)})
|
|
save_state(state_path, state)
|
|
raise failure
|
|
validation = base.validate_cell(entry)
|
|
cell_state["validation"] = validation
|
|
cell_state["status"] = "complete"
|
|
cell_state["completed_at"] = time.time()
|
|
state["completed_cells"] += 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["status"] != "PASS":
|
|
raise RuntimeError("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_base(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, cell in enumerate(ORDER):
|
|
execute_cell(
|
|
index=index,
|
|
cell=cell,
|
|
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)
|
|
print(json.dumps({
|
|
"status": state["status"],
|
|
"completed_cells": state["completed_cells"],
|
|
"gpu_hours_total": state["gpu_hours_total"],
|
|
}, sort_keys=True))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|