Files
aituner/scripts/oracle_gap/run_frontier.py

882 lines
32 KiB
Python

#!/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,
frontier_for_cell,
score_trial,
summarize_trials,
)
SCHEMA = 1
AMENDMENT = "A-OG-4"
AMENDMENT_REASON = (
"close any majority-shifted final boundary to three trials per side "
"without adding rate anchors"
)
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
MAX_CLOSURE_ROUNDS = 3
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, 2.4, 2.5, 2.6, 2.8, 3.0),
}
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")),
"protocol_sha256": sha256_file(
REPO / "docs/opprof/oracle-gap-protocol.md"
),
"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()},
"up_extensions": {
key: list(value) for key, value in UP_EXTENSIONS.items()
},
}
def resume_compatible(old: dict[str, Any], current: dict[str, Any]) -> bool:
immutable = (
"analyzer_sha256",
"p5_client_sha256",
"p3_client_sha256",
"vllm_commit",
"model",
"manifests",
"runtime",
"driver",
"config_details",
"base_rates",
)
if not all(old.get(key) == current.get(key) for key in immutable):
return False
old_extensions = old.get("up_extensions")
current_extensions = current.get("up_extensions")
if old_extensions == current_extensions:
return True
# The only earlier compatible transition was A-OG-1, which appended the
# five registered P06 anchors after the original 2.3-rps list.
if not isinstance(old_extensions, dict) or not isinstance(current_extensions, dict):
return False
return (
old_extensions.get("P01") == current_extensions.get("P01")
and list(old_extensions.get("P06", [])) + [2.4, 2.5, 2.6, 2.8, 3.0]
== current_extensions.get("P06")
)
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):
digest = sha256_file(source)
target = destination / f"{source.stem}.{digest[:12]}{source.suffix}"
if target.exists() and sha256_file(target) != sha256_file(source):
raise RuntimeError(f"content-addressed provenance mismatch: {target}")
if not target.exists():
shutil.copy2(source, target)
fingerprint_path = (
destination / f"fingerprint.{current['repo_commit'][:8]}.json"
)
atomic_json(fingerprint_path, current)
if not (destination / "fingerprint.json").exists():
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 load_cell_scores(phase: str, config: str) -> list[dict[str, Any]]:
paths = sorted((RUN_ROOT / "trials" / f"{phase}-{config}").glob("**/score.json"))
return [json.loads(path.read_text()) for path in paths]
def boundary_closure_needs(rows: list[dict[str, Any]]) -> list[float]:
frontier = frontier_for_cell(rows)
if not frontier["monotone"]:
raise RuntimeError("cannot close a non-monotone frontier")
lower = frontier["lower_feasible_rps"]
upper = frontier["upper_infeasible_rps"]
if lower is None or upper is None:
raise RuntimeError("cannot close an unbracketed frontier")
trials_by_rate = {
float(row["rate_rps"]): int(row["trials"])
for row in frontier["rates"]
}
# Preserve the confirmation convention of running high-to-low.
return [
float(rate)
for rate in (upper, lower)
if trials_by_rate[float(rate)] < 3
]
def run_closure_config(
state: dict[str, Any],
round_index: int,
config: str,
needs: dict[str, list[float]],
) -> None:
stage = f"closure-r{round_index}-{config}"
if state["stages"].get(stage, {}).get("status") == "complete":
return
state["stages"][stage] = {
"status": "starting",
"started_at": time.time(),
"needs": needs,
}
save_state(state)
server = start_server(config, stage, state)
failure = None
try:
for phase in PHASES:
for rate in needs.get(phase, []):
rows = load_cell_scores(phase, config)
repetitions = {
int(row["repetition"])
for row in rows
if float(row["target_rate_rps"]) == rate
}
repetition = 0
while len(repetitions) < 3:
while repetition in repetitions:
repetition += 1
run_trial(
state,
server,
phase,
rate,
repetition,
"boundary-closure",
)
repetitions.add(repetition)
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_boundary_closure(state: dict[str, Any]) -> None:
stage = "boundary-closure"
closure = state["stages"].setdefault(
stage, {"status": "starting", "started_at": time.time(), "rounds": []}
)
save_state(state)
for round_index in range(1, MAX_CLOSURE_ROUNDS + 1):
needs = {
config: {
phase: boundary_closure_needs(load_cell_scores(phase, config))
for phase in PHASES
}
for config in CONFIGS
}
closure["rounds"].append({"round": round_index, "needs": needs})
save_state(state)
if not any(rates for config in needs.values() for rates in config.values()):
closure["status"] = "complete"
closure["completed_at"] = time.time()
save_state(state)
return
for config in CONFIRM_ORDER:
config_needs = {
phase: rates for phase, rates in needs[config].items() if rates
}
if not config_needs:
continue
if float(state["gpu_hours"]) >= GPU_HOUR_LIMIT:
raise RuntimeError("GPU-hour hard cap reached during boundary closure")
run_closure_config(state, round_index, config, config_needs)
remaining = {
config: {
phase: boundary_closure_needs(load_cell_scores(phase, config))
for phase in PHASES
}
for config in CONFIGS
}
closure["final_needs"] = remaining
if not any(rates for config in remaining.values() for rates in config.values()):
closure["status"] = "complete"
closure["completed_at"] = time.time()
save_state(state)
return
raise RuntimeError(
f"failed to close final boundaries in {MAX_CLOSURE_ROUNDS} rounds"
)
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:
if not (
resume
and state.get("status") == "failed"
and resume_compatible(state["fingerprint"], current)
):
raise RuntimeError("resume fingerprint changed incompatibly")
state.setdefault("amendments", []).append(
{
"id": AMENDMENT,
"accepted_at": time.time(),
"reason": AMENDMENT_REASON,
"completed_trials_before": len(
state.get("completed_trials", [])
),
"gpu_hours_before": state.get("gpu_hours"),
"old_fingerprint": state["fingerprint"],
"new_fingerprint": current,
}
)
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)
run_boundary_closure(state)
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,
"max_boundary_closure_rounds": MAX_CLOSURE_ROUNDS,
"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()