Add static-policy oracle gap experiment
This commit is contained in:
671
scripts/oracle_gap/run_frontier.py
Normal file
671
scripts/oracle_gap/run_frontier.py
Normal file
@@ -0,0 +1,671 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Detached, resumable solo-H20 controller for the oracle-gap frontier."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import shlex
|
||||
import shutil
|
||||
import signal
|
||||
import subprocess
|
||||
import time
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from analyze import CONFIGS, PHASES, atomic_json, score_trial, summarize_trials
|
||||
|
||||
|
||||
SCHEMA = 1
|
||||
REMOTE_ROOT = Path("/home/admin/cpfs/wjh/oracle-gap-20260713")
|
||||
RUN_ROOT = REMOTE_ROOT / "runs"
|
||||
STATE = RUN_ROOT / "controller-state.json"
|
||||
PRIVATE = Path("/home/admin/cpfs/wjh/opprof-phase3-private/manifests")
|
||||
MODEL = Path("/home/admin/cpfs/wjh/models/Qwen/Qwen3-30B-A3B")
|
||||
SOURCE = Path("/home/admin/cpfs/wjh/opprof-phase2-dash0-20260711/vllm-v0.24.0")
|
||||
VENV = Path("/tmp/wjh-opprof-phase2-dash0-20260711/.venv")
|
||||
REPO = Path(os.environ.get("AITUNER_ORACLE_REPO", Path(__file__).resolve().parents[2]))
|
||||
P5_CLIENT = REPO / "runs/opprof-phase5/opprof_phase5_client.py"
|
||||
P3_CLIENT_DIR = REPO / "runs/opprof-phase3/provenance"
|
||||
GPU = 0
|
||||
CPU_MASK = "0-19"
|
||||
PORT = 8820
|
||||
GPU_HOUR_LIMIT = 6.0
|
||||
PRIMARY_ORDER = ("C11", "C00", "C01", "C10")
|
||||
CONFIRM_ORDER = tuple(reversed(PRIMARY_ORDER))
|
||||
CONFIG_DETAILS = {
|
||||
"C00": {"mns": 1024, "mbt": 8192, "flags": []},
|
||||
"C10": {"mns": 64, "mbt": 8192, "flags": ["--max-num-seqs", "64"]},
|
||||
"C01": {
|
||||
"mns": 1024,
|
||||
"mbt": 2048,
|
||||
"flags": ["--max-num-batched-tokens", "2048"],
|
||||
},
|
||||
"C11": {
|
||||
"mns": 64,
|
||||
"mbt": 2048,
|
||||
"flags": [
|
||||
"--max-num-seqs", "64", "--max-num-batched-tokens", "2048"
|
||||
],
|
||||
},
|
||||
}
|
||||
BASE_RATES = {
|
||||
"P01": (32.0, 26.0, 36.0, 28.0, 34.0, 30.0),
|
||||
"P06": (1.7, 1.4, 2.0, 1.5, 1.9, 1.6, 1.8),
|
||||
}
|
||||
UP_EXTENSIONS = {"P01": (38.0, 40.0, 42.0), "P06": (2.1, 2.2, 2.3)}
|
||||
DOWN_EXTENSIONS = {"P01": (24.0, 22.0, 20.0), "P06": (1.3, 1.2, 1.1)}
|
||||
TIMELINE = {
|
||||
"P01": {"warmup": 60, "clean": 60, "drain": 120},
|
||||
"P06": {"warmup": 60, "clean": 120, "drain": 240},
|
||||
}
|
||||
|
||||
|
||||
def sha256_file(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as source:
|
||||
for chunk in iter(lambda: source.read(1 << 20), b""):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def run_text(command: list[str], *, check: bool = True) -> str:
|
||||
result = subprocess.run(
|
||||
command, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
|
||||
)
|
||||
if check and result.returncode:
|
||||
raise RuntimeError(
|
||||
f"command failed ({result.returncode}): {shlex.join(command)}\n{result.stdout}"
|
||||
)
|
||||
return result.stdout
|
||||
|
||||
|
||||
def load_state(resume: bool) -> dict[str, Any]:
|
||||
if STATE.exists():
|
||||
if not resume:
|
||||
raise RuntimeError(f"state exists; use --resume: {STATE}")
|
||||
return json.loads(STATE.read_text())
|
||||
return {
|
||||
"schema": SCHEMA,
|
||||
"status": "created",
|
||||
"created_at": time.time(),
|
||||
"gpu_hours": 0.0,
|
||||
"completed_trials": [],
|
||||
"stages": {},
|
||||
"fingerprint": {},
|
||||
"owned_pgids": [],
|
||||
}
|
||||
|
||||
|
||||
def save_state(state: dict[str, Any]) -> None:
|
||||
state["updated_at"] = time.time()
|
||||
state["controller_pid"] = os.getpid()
|
||||
atomic_json(STATE, state)
|
||||
|
||||
|
||||
def compute_apps() -> list[dict[str, Any]]:
|
||||
output = run_text(
|
||||
[
|
||||
"nvidia-smi",
|
||||
"--query-compute-apps=gpu_uuid,pid,process_name,used_memory",
|
||||
"--format=csv,noheader,nounits",
|
||||
],
|
||||
check=False,
|
||||
)
|
||||
rows = []
|
||||
for line in output.splitlines():
|
||||
parts = [part.strip() for part in line.split(",", 3)]
|
||||
if len(parts) == 4 and parts[1].isdigit():
|
||||
rows.append(
|
||||
{
|
||||
"gpu_uuid": parts[0],
|
||||
"pid": int(parts[1]),
|
||||
"process_name": parts[2],
|
||||
"used_memory_mib": int(parts[3].split()[0]),
|
||||
}
|
||||
)
|
||||
return rows
|
||||
|
||||
|
||||
def gpu_snapshot() -> dict[str, Any]:
|
||||
query = run_text(
|
||||
[
|
||||
"nvidia-smi",
|
||||
"--query-gpu=index,name,uuid,memory.used,utilization.gpu,clocks.sm,clocks.mem,power.draw",
|
||||
"--format=csv,noheader,nounits",
|
||||
]
|
||||
)
|
||||
return {
|
||||
"time": time.time(),
|
||||
"gpus": query.splitlines(),
|
||||
"compute_apps": compute_apps(),
|
||||
"loadavg": list(os.getloadavg()),
|
||||
}
|
||||
|
||||
|
||||
def assert_idle() -> None:
|
||||
deadline = time.monotonic() + 30
|
||||
while time.monotonic() < deadline:
|
||||
if not compute_apps():
|
||||
return
|
||||
time.sleep(1)
|
||||
raise RuntimeError(f"GPU host is not idle: {compute_apps()}")
|
||||
|
||||
|
||||
def descendants(root_pid: int) -> set[int]:
|
||||
output = run_text(["ps", "-e", "-o", "pid=,ppid="], check=False)
|
||||
children: dict[int, list[int]] = {}
|
||||
for line in output.splitlines():
|
||||
parts = line.split()
|
||||
if len(parts) == 2:
|
||||
pid, parent = map(int, parts)
|
||||
children.setdefault(parent, []).append(pid)
|
||||
result = {root_pid}
|
||||
pending = [root_pid]
|
||||
while pending:
|
||||
for child in children.get(pending.pop(), []):
|
||||
if child not in result:
|
||||
result.add(child)
|
||||
pending.append(child)
|
||||
return result
|
||||
|
||||
|
||||
def assert_only_server_apps(server_pid: int) -> None:
|
||||
allowed = descendants(server_pid)
|
||||
unexpected = [row for row in compute_apps() if int(row["pid"]) not in allowed]
|
||||
if unexpected:
|
||||
raise RuntimeError(f"unexpected GPU process during run: {unexpected}")
|
||||
|
||||
|
||||
def wait_idle_after_stop() -> None:
|
||||
deadline = time.monotonic() + 60
|
||||
while time.monotonic() < deadline:
|
||||
output = run_text(
|
||||
[
|
||||
"nvidia-smi", "--query-gpu=index,memory.used",
|
||||
"--format=csv,noheader,nounits",
|
||||
]
|
||||
)
|
||||
memory = [int(line.split(",")[1].strip()) for line in output.splitlines()]
|
||||
if not compute_apps() and all(value == 0 for value in memory):
|
||||
return
|
||||
time.sleep(1)
|
||||
raise RuntimeError("GPU memory/processes did not return to zero")
|
||||
|
||||
|
||||
def fingerprint() -> dict[str, Any]:
|
||||
manifests = {}
|
||||
for phase in PHASES:
|
||||
path = PRIVATE / f"{phase}.jsonl"
|
||||
summary = json.loads(path.with_suffix(path.suffix + ".summary.json").read_text())
|
||||
if int(summary["rows"]) != 32768 or summary["sha256"] != sha256_file(path):
|
||||
raise RuntimeError(f"manifest mismatch: {phase}")
|
||||
manifests[phase] = {"path": str(path), "sha256": summary["sha256"]}
|
||||
return {
|
||||
"controller_sha256": sha256_file(Path(__file__).resolve()),
|
||||
"analyzer_sha256": sha256_file(Path(__file__).with_name("analyze.py")),
|
||||
"p5_client_sha256": sha256_file(P5_CLIENT),
|
||||
"p3_client_sha256": sha256_file(P3_CLIENT_DIR / "opprof_phase3_client.py"),
|
||||
"repo_commit": run_text(["git", "-C", str(REPO), "rev-parse", "HEAD"]).strip(),
|
||||
"repo_tree": run_text(["git", "-C", str(REPO), "rev-parse", "HEAD^{tree}"]).strip(),
|
||||
"vllm_commit": run_text(["git", "-C", str(SOURCE), "rev-parse", "HEAD"]).strip(),
|
||||
"model": str(MODEL),
|
||||
"manifests": manifests,
|
||||
"runtime": run_text(
|
||||
[str(VENV / "bin/python"), "-c", "import torch,vllm; print(torch.__version__,torch.version.cuda,vllm.__version__)"]
|
||||
).strip(),
|
||||
"driver": run_text(["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"]).splitlines()[0],
|
||||
"config_details": CONFIG_DETAILS,
|
||||
"base_rates": {key: list(value) for key, value in BASE_RATES.items()},
|
||||
}
|
||||
|
||||
|
||||
def ensure_provenance(current: dict[str, Any]) -> None:
|
||||
destination = RUN_ROOT / "provenance"
|
||||
destination.mkdir(parents=True, exist_ok=True)
|
||||
for source in (
|
||||
Path(__file__).resolve(),
|
||||
Path(__file__).with_name("analyze.py"),
|
||||
P5_CLIENT,
|
||||
P3_CLIENT_DIR / "opprof_phase3_client.py",
|
||||
REPO / "docs/opprof/oracle-gap-protocol.md",
|
||||
):
|
||||
target = destination / source.name
|
||||
if target.exists() and sha256_file(target) != sha256_file(source):
|
||||
raise RuntimeError(f"provenance file changed: {target}")
|
||||
if not target.exists():
|
||||
shutil.copy2(source, target)
|
||||
atomic_json(destination / "fingerprint.json", current)
|
||||
atomic_json(destination / "host-before.json", gpu_snapshot())
|
||||
(destination / "nvidia-smi-q.txt").write_text(run_text(["nvidia-smi", "-q"]))
|
||||
|
||||
|
||||
def rate_label(rate: float) -> str:
|
||||
return f"{rate:.3f}".rstrip("0").rstrip(".")
|
||||
|
||||
|
||||
def trial_key(phase: str, config: str, rate: float, repetition: int) -> str:
|
||||
return f"{phase}-{config}-r{rate_label(rate)}-rep{repetition}"
|
||||
|
||||
|
||||
def derived_manifest(phase: str, rate: float, repetition: int) -> Path:
|
||||
label = f"{phase}-r{rate_label(rate)}-rep{repetition}"
|
||||
output = RUN_ROOT / "manifests" / f"{label}.jsonl"
|
||||
summary_path = output.with_suffix(output.suffix + ".summary.json")
|
||||
source = PRIVATE / f"{phase}.jsonl"
|
||||
domain = int.from_bytes(
|
||||
hashlib.sha256(f"oracle-gap:{label}".encode()).digest()[:4], "big"
|
||||
)
|
||||
if output.exists() and summary_path.exists():
|
||||
summary = json.loads(summary_path.read_text())
|
||||
if summary["sha256"] == sha256_file(output) and summary["source_sha256"] == sha256_file(source):
|
||||
return output
|
||||
raise RuntimeError(f"derived manifest mismatch: {output}")
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
temporary = output.with_name(f"{output.name}.tmp.{os.getpid()}")
|
||||
rows = 0
|
||||
input_sum = 0
|
||||
output_sum = 0
|
||||
with source.open() as src, temporary.open("w") as dst:
|
||||
for line in src:
|
||||
row = json.loads(line)
|
||||
row["token_seed"] = (int(row.get("token_seed", 0)) + domain * 1000003) & ((1 << 63) - 1)
|
||||
row["token_domain"] = domain
|
||||
dst.write(json.dumps(row, sort_keys=True, separators=(",", ":")) + "\n")
|
||||
rows += 1
|
||||
input_sum += int(row["input_tokens"])
|
||||
output_sum += int(row["output_tokens"])
|
||||
dst.flush()
|
||||
os.fsync(dst.fileno())
|
||||
os.replace(temporary, output)
|
||||
summary = {
|
||||
"schema": 1,
|
||||
"phase": phase,
|
||||
"rate": rate,
|
||||
"repetition": repetition,
|
||||
"rows": rows,
|
||||
"input_token_sum": input_sum,
|
||||
"output_token_sum": output_sum,
|
||||
"token_domain": domain,
|
||||
"source_sha256": sha256_file(source),
|
||||
"sha256": sha256_file(output),
|
||||
"invariants": {"rows_32768": rows == 32768, "positive_work": input_sum > 0 and output_sum > 0},
|
||||
}
|
||||
if not all(summary["invariants"].values()):
|
||||
raise RuntimeError(f"derived manifest invalid: {summary}")
|
||||
atomic_json(summary_path, summary)
|
||||
return output
|
||||
|
||||
|
||||
def server_command(config: str) -> list[str]:
|
||||
return [
|
||||
"taskset", "-c", CPU_MASK, str(VENV / "bin/vllm"), "serve", str(MODEL),
|
||||
"--host", "127.0.0.1", "--port", str(PORT),
|
||||
"--tensor-parallel-size", "1", "--enable-chunked-prefill",
|
||||
"--enable-prefix-caching", "--shutdown-timeout", "120",
|
||||
*CONFIG_DETAILS[config]["flags"],
|
||||
]
|
||||
|
||||
|
||||
def wait_ready(process: subprocess.Popen[Any], timeout: float = 300) -> None:
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
if process.poll() is not None:
|
||||
raise RuntimeError(f"server exited before ready: {process.returncode}")
|
||||
try:
|
||||
with urllib.request.urlopen(f"http://127.0.0.1:{PORT}/health", timeout=1) as response:
|
||||
if response.status == 200:
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
time.sleep(1)
|
||||
raise TimeoutError("server readiness timeout")
|
||||
|
||||
|
||||
def validate_startup(log_path: Path, config: str) -> None:
|
||||
log = log_path.read_text(errors="replace")
|
||||
details = CONFIG_DETAILS[config]
|
||||
invariants = {
|
||||
"triton_moe": "Using TRITON Unquantized MoE backend" in log,
|
||||
"tp1": "tensor_parallel_size=1" in log,
|
||||
"mbt": (
|
||||
"Chunked prefill is enabled with max_num_batched_tokens=8192" in log
|
||||
if details["mbt"] == 8192
|
||||
else "'max_num_batched_tokens': 2048" in log
|
||||
),
|
||||
"mns": details["mns"] == 1024 or "'max_num_seqs': 64" in log,
|
||||
}
|
||||
if not all(invariants.values()):
|
||||
raise RuntimeError(f"server startup invariants failed {config}: {invariants}")
|
||||
|
||||
|
||||
def start_server(config: str, stage: str, state: dict[str, Any]) -> dict[str, Any]:
|
||||
assert_idle()
|
||||
directory = RUN_ROOT / "servers" / stage
|
||||
directory.mkdir(parents=True, exist_ok=True)
|
||||
command = server_command(config)
|
||||
echo = (
|
||||
f"RUN_ECHO stage={stage} host=dash0 gpu=0 cpus={CPU_MASK} config={config} "
|
||||
f"model={MODEL} source={SOURCE} manifests={PRIVATE}/P01,P06.jsonl "
|
||||
f"output={RUN_ROOT} expected_server_plus_trials=20-45min budget_cap={GPU_HOUR_LIMIT}H20h"
|
||||
)
|
||||
with (RUN_ROOT / "launch-echo.log").open("a") as handle:
|
||||
handle.write(echo + "\n")
|
||||
print(echo, flush=True)
|
||||
(directory / "command.txt").write_text(shlex.join(command) + "\n")
|
||||
atomic_json(directory / "gpu-before.json", gpu_snapshot())
|
||||
handle = (directory / "server.log").open("ab", buffering=0)
|
||||
environment = os.environ.copy()
|
||||
environment.update(
|
||||
{
|
||||
"CUDA_VISIBLE_DEVICES": str(GPU),
|
||||
"VLLM_OPPROF_DIR": str(directory / "opprof"),
|
||||
"HF_HUB_OFFLINE": "1",
|
||||
"TRANSFORMERS_OFFLINE": "1",
|
||||
"PYTHONUNBUFFERED": "1",
|
||||
"AITUNER_ORACLE_GAP_MARKER": stage,
|
||||
}
|
||||
)
|
||||
process = subprocess.Popen(
|
||||
command, cwd=SOURCE, env=environment, stdout=handle,
|
||||
stderr=subprocess.STDOUT, start_new_session=True,
|
||||
)
|
||||
state["owned_pgids"] = [process.pid]
|
||||
save_state(state)
|
||||
wait_ready(process)
|
||||
validate_startup(directory / "server.log", config)
|
||||
return {
|
||||
"config": config, "stage": stage, "dir": directory, "process": process,
|
||||
"handle": handle, "started_at": time.time(),
|
||||
}
|
||||
|
||||
|
||||
def stop_server(entry: dict[str, Any], state: dict[str, Any]) -> None:
|
||||
process = entry["process"]
|
||||
if process.poll() is None:
|
||||
try:
|
||||
os.kill(process.pid, signal.SIGINT)
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
try:
|
||||
process.wait(timeout=150)
|
||||
except subprocess.TimeoutExpired:
|
||||
for signum, timeout in ((signal.SIGTERM, 10), (signal.SIGKILL, 30)):
|
||||
if process.poll() is not None:
|
||||
break
|
||||
try:
|
||||
os.killpg(process.pid, signum)
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
try:
|
||||
process.wait(timeout=timeout)
|
||||
except subprocess.TimeoutExpired:
|
||||
continue
|
||||
if process.poll() is None:
|
||||
raise TimeoutError(f"server process group did not stop: {entry['stage']}")
|
||||
entry["handle"].close()
|
||||
elapsed = time.time() - float(entry["started_at"])
|
||||
state["gpu_hours"] = float(state["gpu_hours"]) + elapsed / 3600.0
|
||||
state["owned_pgids"] = []
|
||||
atomic_json(entry["dir"] / "gpu-after.json", gpu_snapshot())
|
||||
log = (entry["dir"] / "server.log").read_text(errors="replace")
|
||||
if "mode=drain timeout=120s" not in log:
|
||||
raise RuntimeError(f"server did not use drain shutdown: {entry['stage']}")
|
||||
wait_idle_after_stop()
|
||||
save_state(state)
|
||||
|
||||
|
||||
def client_command(
|
||||
phase: str, rate: float, repetition: int, output: Path
|
||||
) -> list[str]:
|
||||
timeline = TIMELINE[phase]
|
||||
return [
|
||||
"taskset", "-c", CPU_MASK, str(VENV / "bin/python"), str(P5_CLIENT), "run",
|
||||
"--manifest", str(derived_manifest(phase, rate, repetition)),
|
||||
"--base-url", f"http://127.0.0.1:{PORT}", "--model", str(MODEL),
|
||||
"--load-point", "moderate", "--fixed-request-rate", str(rate),
|
||||
"--max-concurrency", "256", "--ignore-eos", "--temperature", "0",
|
||||
"--warmup-seconds", str(timeline["warmup"]),
|
||||
"--clean-segment-seconds", str(timeline["clean"]),
|
||||
"--num-clean-segments", "1", "--post-clean-seconds", "0",
|
||||
"--drain-timeout-seconds", str(timeline["drain"]),
|
||||
"--workload-seed", "20260712", "--server-seed", "20260712",
|
||||
"--result-dir", str(output / "client"),
|
||||
]
|
||||
|
||||
|
||||
def run_trial(
|
||||
state: dict[str, Any], server: dict[str, Any], phase: str, rate: float,
|
||||
repetition: int, role: str,
|
||||
) -> dict[str, Any]:
|
||||
config = str(server["config"])
|
||||
key = trial_key(phase, config, rate, repetition)
|
||||
output = RUN_ROOT / "trials" / f"{phase}-{config}" / f"rate-{rate_label(rate)}" / f"rep-{repetition}"
|
||||
score_path = output / "score.json"
|
||||
if key in state["completed_trials"]:
|
||||
if not score_path.exists():
|
||||
raise RuntimeError(f"completed state lacks score: {key}")
|
||||
return json.loads(score_path.read_text())
|
||||
output.mkdir(parents=True, exist_ok=True)
|
||||
command = client_command(phase, rate, repetition, output)
|
||||
(output / "command.txt").write_text(shlex.join(command) + "\n")
|
||||
if server["process"].poll() is not None:
|
||||
raise RuntimeError(f"server exited before trial: {key}")
|
||||
assert_only_server_apps(int(server["process"].pid))
|
||||
environment = os.environ.copy()
|
||||
environment["PYTHONPATH"] = str(P3_CLIENT_DIR) + os.pathsep + environment.get("PYTHONPATH", "")
|
||||
started = time.time()
|
||||
with (output / "client.log").open("ab", buffering=0) as handle:
|
||||
result = subprocess.run(
|
||||
command, cwd=REPO, env=environment, stdout=handle,
|
||||
stderr=subprocess.STDOUT, timeout=900,
|
||||
)
|
||||
assert_only_server_apps(int(server["process"].pid))
|
||||
if server["process"].poll() is not None:
|
||||
raise RuntimeError(f"server exited during trial: {key}")
|
||||
client_result = output / "client/result.json"
|
||||
client_sanity = output / "client/sanity.json"
|
||||
if not client_result.exists() or not client_sanity.exists():
|
||||
raise RuntimeError(f"client produced no result: {key} rc={result.returncode}")
|
||||
sanity = json.loads(client_sanity.read_text())["invariants"]
|
||||
failed = [name for name, passed in sanity.items() if not passed]
|
||||
allowed = {"moderate_offered_within_5pct"}
|
||||
if result.returncode != 0 and not failed:
|
||||
raise RuntimeError(f"client failed without sanity marker: {key}")
|
||||
if set(failed) - allowed:
|
||||
raise RuntimeError(f"client validity failure {key}: {failed}")
|
||||
score = score_trial(
|
||||
output / "client/requests.jsonl", client_result, phase=phase,
|
||||
config=config, target_rate=rate, repetition=repetition, role=role,
|
||||
)
|
||||
score["wall_seconds"] = time.time() - started
|
||||
score["client_returncode"] = result.returncode
|
||||
score["client_failed_invariants"] = failed
|
||||
score["manifest_sha256"] = sha256_file(derived_manifest(phase, rate, repetition))
|
||||
atomic_json(score_path, score)
|
||||
state["completed_trials"].append(key)
|
||||
state["last_trial"] = {
|
||||
"key": key, "feasible": score["feasible"], "pass_rate": score["pass_rate"],
|
||||
"goodput": score["slo_goodput_rps"], "completed_at": time.time(),
|
||||
}
|
||||
save_state(state)
|
||||
print(
|
||||
f"TRIAL {key} feasible={score['feasible']} pass={score['pass_rate']:.6f} "
|
||||
f"goodput={score['slo_goodput_rps']:.6f} lagmax={score['schedule_lag_ms']['max']:.3f}",
|
||||
flush=True,
|
||||
)
|
||||
return score
|
||||
|
||||
|
||||
def primary_bracket(rows: list[dict[str, Any]]) -> tuple[float, float] | None:
|
||||
verdict = {float(row["target_rate_rps"]): bool(row["feasible"]) for row in rows}
|
||||
ordered = sorted(verdict)
|
||||
passes = [rate for rate in ordered if verdict[rate]]
|
||||
failures = [rate for rate in ordered if not verdict[rate]]
|
||||
if not passes or not failures:
|
||||
return None
|
||||
lower = max(passes)
|
||||
upper = min((rate for rate in failures if rate > lower), default=None)
|
||||
if upper is None or any(verdict[rate] for rate in ordered if rate > upper):
|
||||
raise RuntimeError(f"non-monotone primary frontier: {verdict}")
|
||||
return lower, upper
|
||||
|
||||
|
||||
def run_primary_config(state: dict[str, Any], config: str) -> None:
|
||||
stage = f"primary-{config}"
|
||||
if state["stages"].get(stage, {}).get("status") == "complete":
|
||||
return
|
||||
state["stages"][stage] = {"status": "starting", "started_at": time.time()}
|
||||
save_state(state)
|
||||
server = start_server(config, stage, state)
|
||||
failure = None
|
||||
try:
|
||||
for phase in PHASES:
|
||||
phase_rows = [run_trial(state, server, phase, rate, 0, "primary") for rate in BASE_RATES[phase]]
|
||||
bracket = primary_bracket(phase_rows)
|
||||
if bracket is None:
|
||||
extension = UP_EXTENSIONS[phase] if all(row["feasible"] for row in phase_rows) else DOWN_EXTENSIONS[phase]
|
||||
for rate in extension:
|
||||
phase_rows.append(run_trial(state, server, phase, rate, 0, "primary-extension"))
|
||||
bracket = primary_bracket(phase_rows)
|
||||
if bracket is not None:
|
||||
break
|
||||
if bracket is None:
|
||||
raise RuntimeError(f"failed to bracket {phase}-{config}")
|
||||
state["stages"][stage].setdefault("brackets", {})[phase] = list(bracket)
|
||||
save_state(state)
|
||||
except Exception as error:
|
||||
failure = error
|
||||
finally:
|
||||
try:
|
||||
stop_server(server, state)
|
||||
except Exception as error:
|
||||
failure = failure or error
|
||||
if failure is not None:
|
||||
state["stages"][stage]["status"] = "failed"
|
||||
state["stages"][stage]["failure"] = repr(failure)
|
||||
state["status"] = "failed"
|
||||
save_state(state)
|
||||
raise failure
|
||||
state["stages"][stage]["status"] = "complete"
|
||||
state["stages"][stage]["completed_at"] = time.time()
|
||||
save_state(state)
|
||||
|
||||
|
||||
def run_confirm_config(state: dict[str, Any], config: str) -> None:
|
||||
stage = f"confirm-{config}"
|
||||
if state["stages"].get(stage, {}).get("status") == "complete":
|
||||
return
|
||||
brackets = state["stages"][f"primary-{config}"]["brackets"]
|
||||
state["stages"][stage] = {"status": "starting", "started_at": time.time(), "brackets": brackets}
|
||||
save_state(state)
|
||||
server = start_server(config, stage, state)
|
||||
failure = None
|
||||
try:
|
||||
for phase in PHASES:
|
||||
lower, upper = (float(value) for value in brackets[phase])
|
||||
for rate in (upper, lower):
|
||||
for repetition in (1, 2):
|
||||
run_trial(state, server, phase, rate, repetition, "boundary-confirmation")
|
||||
except Exception as error:
|
||||
failure = error
|
||||
finally:
|
||||
try:
|
||||
stop_server(server, state)
|
||||
except Exception as error:
|
||||
failure = failure or error
|
||||
if failure is not None:
|
||||
state["stages"][stage]["status"] = "failed"
|
||||
state["stages"][stage]["failure"] = repr(failure)
|
||||
state["status"] = "failed"
|
||||
save_state(state)
|
||||
raise failure
|
||||
state["stages"][stage]["status"] = "complete"
|
||||
state["stages"][stage]["completed_at"] = time.time()
|
||||
save_state(state)
|
||||
|
||||
|
||||
def cleanup_recorded(state: dict[str, Any]) -> None:
|
||||
for pgid in state.get("owned_pgids", []):
|
||||
try:
|
||||
os.killpg(int(pgid), signal.SIGKILL)
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
state["owned_pgids"] = []
|
||||
save_state(state)
|
||||
wait_idle_after_stop()
|
||||
|
||||
|
||||
def execute(resume: bool) -> None:
|
||||
RUN_ROOT.mkdir(parents=True, exist_ok=True)
|
||||
state = load_state(resume)
|
||||
if resume and state.get("owned_pgids"):
|
||||
cleanup_recorded(state)
|
||||
assert_idle()
|
||||
current = fingerprint()
|
||||
if state["fingerprint"] and state["fingerprint"] != current:
|
||||
raise RuntimeError("resume fingerprint changed")
|
||||
state["fingerprint"] = current
|
||||
state["status"] = "running"
|
||||
save_state(state)
|
||||
ensure_provenance(current)
|
||||
for config in PRIMARY_ORDER:
|
||||
if float(state["gpu_hours"]) >= GPU_HOUR_LIMIT:
|
||||
raise RuntimeError("GPU-hour hard cap reached before primary completion")
|
||||
run_primary_config(state, config)
|
||||
for config in CONFIRM_ORDER:
|
||||
if float(state["gpu_hours"]) >= GPU_HOUR_LIMIT:
|
||||
raise RuntimeError("GPU-hour hard cap reached before confirmations")
|
||||
run_confirm_config(state, config)
|
||||
score_paths = sorted((RUN_ROOT / "trials").glob("**/score.json"))
|
||||
summary = summarize_trials([json.loads(path.read_text()) for path in score_paths])
|
||||
summary["gpu_hours"] = state["gpu_hours"]
|
||||
summary["trial_files"] = len(score_paths)
|
||||
atomic_json(RUN_ROOT / "metrics.json", summary)
|
||||
state["status"] = "complete"
|
||||
state["completed_at"] = time.time()
|
||||
state["verdict"] = summary["verdict"]
|
||||
save_state(state)
|
||||
print(json.dumps({"status": "complete", "verdict": summary["verdict"], "gpu_hours": state["gpu_hours"]}, sort_keys=True))
|
||||
|
||||
|
||||
def plan() -> dict[str, Any]:
|
||||
primary = sum(len(BASE_RATES[phase]) for phase in PHASES) * len(CONFIGS)
|
||||
confirmations = 2 * 2 * len(PHASES) * len(CONFIGS)
|
||||
return {
|
||||
"schema": 1,
|
||||
"placement": "serialized solo GPU0",
|
||||
"primary_order": list(PRIMARY_ORDER),
|
||||
"confirm_order": list(CONFIRM_ORDER),
|
||||
"primary_trials_without_extensions": primary,
|
||||
"confirmation_trials": confirmations,
|
||||
"expected_total_trials": primary + confirmations,
|
||||
"expected_h20_hours": "3.0-4.0",
|
||||
"hard_cap_h20_hours": GPU_HOUR_LIMIT,
|
||||
"rates": {key: list(value) for key, value in BASE_RATES.items()},
|
||||
"timelines": TIMELINE,
|
||||
"output": str(RUN_ROOT),
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
sub = parser.add_subparsers(dest="command", required=True)
|
||||
run = sub.add_parser("run")
|
||||
run.add_argument("--resume", action="store_true")
|
||||
sub.add_parser("plan")
|
||||
sub.add_parser("status")
|
||||
args = parser.parse_args()
|
||||
if args.command == "run":
|
||||
execute(args.resume)
|
||||
elif args.command == "plan":
|
||||
print(json.dumps(plan(), indent=2, sort_keys=True))
|
||||
else:
|
||||
print(STATE.read_text() if STATE.exists() else json.dumps({"status": "not-started"}))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user