aituner/runs/opprof-phase3/provenance/opprof_phase3_matrix.py

#!/usr/bin/env python3
"""Detached, resumable four-GPU controller for the Phase-3 matrix."""

from __future__ import annotations

import argparse
import gzip
import hashlib
import json
import math
import os
import re
import shlex
import shutil
import signal
import subprocess
import time
import urllib.request
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import opprof_phase3_controller as common

SCHEMA = 1
WORKDIR = Path("/home/admin/cpfs/wjh/opprof-phase3-dash0-20260712")
RUN_ROOT = WORKDIR / "runs/phase3"
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")
CLIENT = WORKDIR / "scripts/opprof_phase3_client.py"
STATE = RUN_ROOT / "controller-state.json"
GPU_LIMIT = 4
GPU_HOUR_LIMIT = 16.0
PRIOR_GPU_HOURS = 3.964 + 427.43714332580566 / 3600
CPU_MAP = common.CPU_MAP
PROFILE_CONFIG = {
    "profiler": "torch",
    "torch_profiler_with_stack": True,
    "torch_profiler_record_shapes": True,
    "torch_profiler_use_gzip": True,
    "ignore_frontend": True,
    "wait_iterations": 0,
    "warmup_iterations": 2,
    "active_iterations": 8,
}
CONFIGS = {
    "C00": {"tp": 1, "mns": 1024, "mbt": 8192, "flags": []},
    "C10": {
        "tp": 1,
        "mns": 64,
        "mbt": 8192,
        "flags": ["--max-num-seqs", "64"],
    },
    "C01": {
        "tp": 1,
        "mns": 1024,
        "mbt": 2048,
        "flags": ["--max-num-batched-tokens", "2048"],
    },
    "C11": {
        "tp": 1,
        "mns": 64,
        "mbt": 2048,
        "flags": [
            "--max-num-seqs",
            "64",
            "--max-num-batched-tokens",
            "2048",
        ],
    },
    "C00-TP2": {"tp": 2, "mns": 1024, "mbt": 8192, "flags": []},
}
LONG_BURST = {"P04", "P06", "P07", "P08", "P11"}
SENTINELS = ("P01", "P03", "P06", "P10")
PATTERNS = tuple(f"P{index:02d}" for index in range(1, 12))


@dataclass(frozen=True)
class Cell:
    pattern: str
    config: str

    @property
    def cell_id(self) -> str:
        return f"{self.pattern}-{self.config}"

    @property
    def width(self) -> int:
        return int(CONFIGS[self.config]["tp"])


@dataclass(frozen=True)
class Assignment:
    cell: Cell
    gpus: tuple[int, ...]


def drain_budget(pattern: str) -> int:
    if pattern == "P10":
        return 600
    if pattern in LONG_BURST:
        return 240
    return 120


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 atomic_json(path: Path, value: Any) -> None:
    common.atomic_json(path, value)


def run_text(command: list[str], check: bool = True) -> str:
    return common.run_text(command, check=check)


def numeric_sanity(values: list[float | int]) -> dict[str, Any]:
    finite = [float(value) for value in values if math.isfinite(float(value))]
    return {
        "n": len(values),
        "finite_n": len(finite),
        "missing_n": len(values) - len(finite),
        "min": min(finite) if finite else None,
        "max": max(finite) if finite else None,
        "distinct_n": len(set(finite)),
    }


def cells() -> list[Cell]:
    result = [Cell(pattern, "C00") for pattern in PATTERNS]
    for pattern in SENTINELS:
        result.extend(Cell(pattern, config) for config in ("C10", "C01", "C11"))
    result.append(Cell("P10", "C00-TP2"))
    return sorted(
        result,
        key=lambda cell: hashlib.sha256(
            f"20260713:{cell.pattern}:{cell.config}".encode()
        ).hexdigest(),
    )


def pack_cells(items: list[Cell]) -> list[list[Assignment]]:
    waves: list[list[Assignment]] = []
    current: list[Assignment] = []
    slot = 0
    for cell in items:
        if slot + cell.width > GPU_LIMIT:
            waves.append(current)
            current = []
            slot = 0
        current.append(Assignment(cell, tuple(range(slot, slot + cell.width))))
        slot += cell.width
    if current:
        waves.append(current)
    assert sum(len(wave) for wave in waves) == len(items)
    assert all(sum(item.cell.width for item in wave) <= GPU_LIMIT for wave in waves)
    return waves


def load_state(resume: bool) -> dict[str, Any]:
    if STATE.exists():
        if not resume:
            raise RuntimeError("controller state exists; use --resume")
        return json.loads(STATE.read_text())
    return {
        "schema": SCHEMA,
        "status": "created",
        "created_at": time.time(),
        "controller_pid": os.getpid(),
        "gpu_hours_total": PRIOR_GPU_HOURS,
        "gpu_hours_this_stage": 0.0,
        "completed_measured_runs": 0,
        "completed_burnins": 0,
        "drain_quarantined_runs": 0,
        "clean_window_failures": 0,
        "missing_trace_files": 0,
        "stages": {},
        "fingerprint": {},
    }


def save_state(state: dict[str, Any]) -> None:
    state["controller_pid"] = os.getpid()
    state["updated_at"] = time.time()
    atomic_json(STATE, state)


def manifest_fingerprint() -> dict[str, Any]:
    result = {}
    for pattern in PATTERNS:
        path = PRIVATE / f"{pattern}.jsonl"
        summary_path = path.with_suffix(path.suffix + ".summary.json")
        if not path.exists() or not summary_path.exists():
            raise RuntimeError(f"manifest missing: {path}")
        summary = json.loads(summary_path.read_text())
        expected_rows = 4011 if pattern == "P10" else 32768
        if summary["rows"] != expected_rows or summary["sha256"] != sha256_file(path):
            raise RuntimeError(f"manifest verification failed: {pattern}")
        result[pattern] = {
            "path": str(path),
            "rows": summary["rows"],
            "sha256": summary["sha256"],
        }
    return result


def fingerprint() -> dict[str, Any]:
    return {
        "source_commit": run_text(
            ["git", "-C", str(SOURCE), "rev-parse", "HEAD"]
        ).strip(),
        "source_tree": run_text(
            ["git", "-C", str(SOURCE), "rev-parse", "HEAD^{tree}"]
        ).strip(),
        "client_sha256": sha256_file(CLIENT),
        "controller_sha256": sha256_file(Path(__file__)),
        "common_controller_sha256": sha256_file(
            Path(common.__file__).resolve()
        ),
        "model": str(MODEL),
        "cpu_map": {str(key): value for key, value in CPU_MAP.items()},
        "manifests": manifest_fingerprint(),
        "cells": [cell.cell_id for cell in cells()],
    }


def ensure_provenance() -> None:
    destination = RUN_ROOT / "provenance"
    destination.mkdir(parents=True, exist_ok=True)
    sources = [CLIENT, Path(__file__).resolve(), Path(common.__file__).resolve()]
    checksums = {}
    for source in sources:
        target = destination / source.name
        digest = sha256_file(source)
        if target.exists() and sha256_file(target) != digest:
            target = destination / f"{source.stem}.{digest[:12]}{source.suffix}"
        if target.exists() and sha256_file(target) != digest:
            raise RuntimeError(f"content-addressed provenance mismatch: {target}")
        if not target.exists():
            shutil.copy2(source, target)
        checksums[target.name] = digest
    patch_dir = WORKDIR / "provenance/patches-vllm-0.24.0-opprof"
    checksums["patches"] = {
        path.name: sha256_file(path) for path in sorted(patch_dir.glob("0*.patch"))
    }
    atomic_json(destination / "sha256.json", checksums)


def preflight(stage_dir: Path, gpus: list[int]) -> None:
    stage_dir.mkdir(parents=True, exist_ok=True)
    common.preflight(gpus, stage_dir)
    free_kb = shutil.disk_usage(RUN_ROOT.parent).free
    if free_kb < 100 * (1 << 30):
        raise RuntimeError("CPFS free space below 100 GiB")


def cpu_mask(gpus: tuple[int, ...]) -> str:
    ranges = [CPU_MAP[gpu] for gpu in gpus]
    return ",".join(ranges)


def server_command(
    assignment: Assignment, port: int, trace_dir: Path
) -> list[str]:
    config = {
        **PROFILE_CONFIG,
        "torch_profiler_dir": str(trace_dir),
    }
    details = CONFIGS[assignment.cell.config]
    return [
        "taskset",
        "-c",
        cpu_mask(assignment.gpus),
        str(VENV / "bin/vllm"),
        "serve",
        str(MODEL),
        "--host",
        "127.0.0.1",
        "--port",
        str(port),
        "--tensor-parallel-size",
        str(details["tp"]),
        "--enable-chunked-prefill",
        "--enable-prefix-caching",
        "--shutdown-timeout",
        "120",
        "--profiler-config",
        json.dumps(config, separators=(",", ":")),
        *details["flags"],
    ]


def client_command(
    assignment: Assignment,
    port: int,
    run_dir: Path,
    load_point: str,
    profile: bool,
    burnin: bool,
    saturation_result: Path | None,
) -> list[str]:
    if burnin:
        warmup, segment, segments = 0, 20, 3
    else:
        warmup, segment, segments = 60, 80, 3
    drain = drain_budget(assignment.cell.pattern)
    command = [
        "taskset",
        "-c",
        cpu_mask(assignment.gpus),
        str(VENV / "bin/python"),
        str(CLIENT),
        "run",
        "--manifest",
        str(PRIVATE / f"{assignment.cell.pattern}.jsonl"),
        "--base-url",
        f"http://127.0.0.1:{port}",
        "--model",
        str(MODEL),
        "--load-point",
        load_point,
        "--max-concurrency",
        "256",
        "--ignore-eos",
        "--temperature",
        "0",
        "--warmup-seconds",
        str(warmup),
        "--clean-segment-seconds",
        str(segment),
        "--num-clean-segments",
        str(segments),
        "--recovery-seconds",
        "30",
        "--drain-timeout-seconds",
        str(drain),
        "--workload-seed",
        "20260712",
        "--server-seed",
        "20260712",
        "--result-dir",
        str(run_dir / "client"),
    ]
    if load_point == "saturation":
        command.extend(("--request-rate", "inf"))
    else:
        if saturation_result is None:
            raise RuntimeError("moderate run lacks saturation result")
        command.extend(
            (
                "--saturation-result",
                str(saturation_result),
                "--rate-fraction",
                "0.60",
            )
        )
    if profile:
        command.extend(
            (
                "--profile-after-clean",
                "--num-profile-windows",
                "2",
                "--profile-warmup-iterations",
                "2",
                "--profile-active-iterations",
                "8",
                "--profile-trace-dir",
                str(run_dir / "trace-tmp"),
                "--profile-timeout-seconds",
                "120",
            )
        )
    return command


def wait_ready(entries: list[dict[str, Any]]) -> None:
    pending = {entry["port"]: entry for entry in entries}
    deadline = time.monotonic() + 300
    while pending and time.monotonic() < deadline:
        for port, entry in list(pending.items()):
            if entry["server"].poll() is not None:
                raise RuntimeError(f"server exited before ready: {entry['run_id']}")
            try:
                with urllib.request.urlopen(
                    f"http://127.0.0.1:{port}/health", timeout=1
                ) as response:
                    if response.status == 200:
                        del pending[port]
            except Exception:
                pass
        time.sleep(1)
    if pending:
        raise TimeoutError(f"server readiness timeout: {sorted(pending)}")


def stop_servers(entries: list[dict[str, Any]]) -> None:
    live = [entry for entry in entries if entry["server"].poll() is None]
    for entry in live:
        try:
            os.kill(entry["server"].pid, signal.SIGINT)
        except ProcessLookupError:
            pass
    deadline = time.monotonic() + 150
    while time.monotonic() < deadline and any(
        entry["server"].poll() is None for entry in live
    ):
        time.sleep(1)
    remaining = [entry for entry in live if entry["server"].poll() is None]
    for signum in (signal.SIGTERM, signal.SIGKILL):
        if not remaining:
            break
        for entry in remaining:
            try:
                os.killpg(entry["server"].pid, signum)
            except ProcessLookupError:
                pass
        time.sleep(5)
        remaining = [entry for entry in remaining if entry["server"].poll() is None]
    if remaining:
        raise RuntimeError("server process groups survived SIGKILL")


def trace_summary(path: Path) -> dict[str, Any]:
    opener = gzip.open if path.suffix == ".gz" else open
    with opener(path, "rt", encoding="utf-8") as source:
        payload = json.load(source)
    durations: dict[str, float] = {}
    kernel_count = 0
    steps = set()
    executes = 0
    for event in payload.get("traceEvents", []):
        name = str(event.get("name", ""))
        if event.get("cat") == "user_annotation":
            match = re.fullmatch(r"ProfilerStep#(\d+)", name)
            if match:
                steps.add(int(match.group(1)))
            if name.startswith("execute_"):
                executes += 1
        if event.get("cat") == "kernel":
            duration = float(event.get("dur", 0))
            if duration < 0:
                raise RuntimeError(f"negative kernel duration in {path}")
            family = common.classify_kernel(name)
            durations[family] = durations.get(family, 0.0) + duration
            kernel_count += 1
    total = sum(durations.values())
    shares = {
        family: duration / total for family, duration in sorted(durations.items())
    }
    classifiable = 1.0 - shares.get("other", 0.0)
    valid = (
        kernel_count > 0
        and steps == set(range(2, 10))
        and executes == 8
        and classifiable >= 0.70
    )
    return {
        "path": str(path),
        "sha256": sha256_file(path),
        "bytes": path.stat().st_size,
        "kernel_count": kernel_count,
        "profiler_steps": sorted(steps),
        "execute_annotations": executes,
        "duration_us": durations,
        "shares": shares,
        "classifiable_fraction": classifiable,
        "valid": valid,
    }


def validate_layer1(run_dir: Path) -> dict[str, Any]:
    streams = sorted((run_dir / "opprof").glob("*.jsonl"))
    sidecars = sorted((run_dir / "opprof").glob("*.jsonl.footer.json"))
    if len(streams) != 1 or len(sidecars) != 1:
        raise RuntimeError(
            f"{run_dir}: expected one Layer-1 stream/sidecar, "
            f"got {len(streams)}/{len(sidecars)}"
        )
    raw = streams[0].read_bytes()
    if not raw.endswith(b"\n"):
        raise RuntimeError(f"partial Layer-1 line: {streams[0]}")
    decoded = [json.loads(line) for line in raw.splitlines()]
    if not decoded or decoded[-1].get("record_type") != "footer":
        raise RuntimeError(f"clean-close footer missing: {streams[0]}")
    records, footer = decoded[:-1], decoded[-1]
    sidecar = json.loads(sidecars[0].read_text())
    indices = [int(record["step_index"]) for record in records]
    invariants = {
        "schema_1": all(item.get("schema") == 1 for item in decoded)
        and sidecar.get("schema") == 1,
        "steps_unique_contiguous": sorted(indices) == list(range(len(indices))),
        "footer_written_matches": footer["written_records"] == len(records),
        "footer_balanced": footer["encoded_records"]
        == footer["written_records"] + footer["dropped_records"],
        "zero_drops": footer["dropped_records"] == 0
        and all(record["dropped_records_before"] == 0 for record in records),
        "sidecar_final": sidecar.get("final") is True,
        "sidecar_agrees": all(
            sidecar[counter] == footer[counter]
            for counter in ("encoded_records", "written_records", "dropped_records")
        ),
        "token_composition": all(
            record["prefill_tokens"] + record["decode_tokens"]
            >= 0
            for record in records
        ),
        "cudagraph_identity": all(
            record["cudagraph"]["bucket_tokens"]
            == record["cudagraph"]["unpadded_tokens"]
            + record["cudagraph"]["padding_tokens"]
            for record in records
        ),
    }
    if not all(invariants.values()):
        raise RuntimeError(f"Layer-1 invariant failure: {run_dir}: {invariants}")
    return {
        "stream": str(streams[0]),
        "stream_sha256": sha256_file(streams[0]),
        "sidecar": str(sidecars[0]),
        "records": len(records),
        "last_step_index": max(indices),
        "footer": footer,
        "invariants": invariants,
        "numeric": {
            "prefill_tokens": numeric_sanity(
                [record["prefill_tokens"] for record in records]
            ),
            "decode_tokens": numeric_sanity(
                [record["decode_tokens"] for record in records]
            ),
            "scheduled_requests": numeric_sanity(
                [record["scheduled_requests"] for record in records]
            ),
            "padding_tokens": numeric_sanity(
                [record["cudagraph"]["padding_tokens"] for record in records]
            ),
        },
    }


def summarize_request_failures(
    requests: list[dict[str, Any]], clean_start: float, clean_end: float
) -> dict[str, Any]:
    failed = [request for request in requests if not request["success"]]
    clean_failed = [
        request
        for request in failed
        if clean_start <= float(request["completed_s"]) < clean_end
    ]
    excluded_kinds: dict[str, int] = {}
    for request in failed:
        if clean_start <= float(request["completed_s"]) < clean_end:
            continue
        key = str(request.get("error_kind") or "unknown")
        excluded_kinds[key] = excluded_kinds.get(key, 0) + 1
    return {
        "failed": len(failed),
        "clean_failed": len(clean_failed),
        "excluded": len(failed) - len(clean_failed),
        "excluded_kinds": excluded_kinds,
    }


def p10_warmup_stability(run_dir: Path, t0_mono_ns: int) -> dict[str, Any]:
    streams = sorted((run_dir / "opprof").glob("*.jsonl"))
    if len(streams) != 1:
        return {
            "passed": False,
            "reason": f"expected one Layer-1 stream, found {len(streams)}",
            "step_counts": [0, 0, 0],
            "scheduled_token_throughput": [None, None, None],
            "mean_scheduled_token_throughput": None,
            "slope_tokens_per_second_squared": None,
            "normalized_drift": None,
        }
    records = []
    try:
        for line in streams[0].read_text().splitlines():
            item = json.loads(line)
            if "step_index" in item:
                records.append(item)
    except (OSError, ValueError) as error:
        return {
            "passed": False,
            "reason": f"Layer-1 decode failed: {error}",
            "step_counts": [0, 0, 0],
            "scheduled_token_throughput": [None, None, None],
            "mean_scheduled_token_throughput": None,
            "slope_tokens_per_second_squared": None,
            "normalized_drift": None,
        }
    indices = [int(item["step_index"]) for item in records]
    continuous = indices == list(range(len(indices)))
    step_counts = [0, 0, 0]
    token_counts = [0, 0, 0]
    for item in records:
        if not item.get("model_executed"):
            continue
        relative_s = (int(item["submit_mono_ns"]) - t0_mono_ns) / 1e9
        if not 45 <= relative_s < 60:
            continue
        bin_index = min(2, int((relative_s - 45) // 5))
        step_counts[bin_index] += 1
        token_counts[bin_index] += int(item["prefill_tokens"]) + int(
            item["decode_tokens"]
        )
    rates = [tokens / 5.0 for tokens in token_counts]
    mean_rate = sum(rates) / len(rates)
    midpoints = [47.5, 52.5, 57.5]
    mean_midpoint = sum(midpoints) / len(midpoints)
    slope = sum(
        (point - mean_midpoint) * (rate - mean_rate)
        for point, rate in zip(midpoints, rates, strict=True)
    ) / sum((point - mean_midpoint) ** 2 for point in midpoints)
    normalized_drift = (
        abs(slope) * 15.0 / mean_rate if mean_rate > 0 else math.inf
    )
    finite_positive = all(math.isfinite(rate) and rate > 0 for rate in rates)
    passed = (
        continuous
        and all(count >= 16 for count in step_counts)
        and finite_positive
        and math.isfinite(normalized_drift)
        and normalized_drift <= 0.10
    )
    return {
        "passed": passed,
        "reason": None if passed else "A-P3-6 stabilization criterion not met",
        "window_seconds": [45.0, 60.0],
        "bin_seconds": 5.0,
        "step_counts": step_counts,
        "scheduled_tokens": token_counts,
        "scheduled_token_throughput": rates,
        "mean_scheduled_token_throughput": mean_rate,
        "slope_tokens_per_second_squared": slope,
        "normalized_drift": normalized_drift,
        "normalized_drift_limit": 0.10,
        "step_indices_continuous": continuous,
    }


def validate_client(
    run_dir: Path, pattern: str, profile: bool, burnin: bool
) -> dict[str, Any]:
    result = json.loads((run_dir / "client/result.json").read_text())
    sanity = json.loads((run_dir / "client/sanity.json").read_text())
    failed = [
        key
        for key, value in sanity["invariants"].items()
        if not value and key != "drain_within_timeout"
    ]
    quarantined = float(result["drain_seconds"]) > drain_budget(pattern)
    expected_duration = 60 if burnin else 240
    warmup_completions = 0
    requests = []
    for line in (run_dir / "client/requests.jsonl").read_text().splitlines():
        request = json.loads(line)
        requests.append(request)
        if 0 <= request["completed_s"] < result["warmup_seconds"] and request["success"]:
            warmup_completions += 1
    clean_start = float(result["clean"]["start_s"])
    clean_end = float(result["clean"]["end_s"])
    failure_summary = summarize_request_failures(requests, clean_start, clean_end)
    request_rate = result["request_rate"]
    warmup_required = 32
    if pattern != "P10" and request_rate != "inf":
        warmup_required = min(
            32,
            max(1, math.floor(float(request_rate) * result["warmup_seconds"])),
        )
    warmup_stability = (
        p10_warmup_stability(run_dir, int(result["t0_mono_ns"]))
        if pattern == "P10" and not burnin
        else None
    )
    if burnin:
        warmup_passed = True
        warmup_gate_branch = "burnin"
    elif pattern == "P10" and warmup_completions >= 32:
        warmup_passed = True
        warmup_gate_branch = "count"
    elif (
        pattern == "P10"
        and warmup_completions >= 16
        and warmup_stability is not None
        and warmup_stability["passed"]
    ):
        warmup_passed = True
        warmup_gate_branch = "telemetry"
    else:
        warmup_passed = warmup_completions >= warmup_required
        warmup_gate_branch = "count" if warmup_passed else "failed"
    invariants = {
        "client_sanity": not failed,
        "clean_duration": math.isclose(
            float(result["clean"]["duration_s"]), expected_duration
        ),
        "clean_failures_zero": result["clean"]["failed"] == 0
        and failure_summary["clean_failed"] == 0,
        "failed_records_accounted": result["failed_records"]
        == failure_summary["failed"],
        "manifest_no_wrap": not result["manifest_wrapped"]
        and not result["manifest_exhausted"],
        "warmup_completions": warmup_passed,
        "profile_count": len(result["profiles"]) == (2 if profile else 0),
        "profile_after_clean": all(
            item["start_call_s"] >= result["warmup_seconds"] + expected_duration
            for item in result["profiles"]
        ),
        "drain_re_adjudicated": not quarantined,
    }
    non_drain_invariants = {
        key: value for key, value in invariants.items() if key != "drain_re_adjudicated"
    }
    if not all(non_drain_invariants.values()):
        raise RuntimeError(
            f"client invariant failure: {run_dir}: {invariants}; failed={failed}; "
            f"warmup_completions={warmup_completions}; "
            f"warmup_gate_branch={warmup_gate_branch}; "
            f"warmup_stability={warmup_stability}"
        )
    return {
        "result": result,
        "sanity": sanity,
        "request_count": len(requests),
        "warmup_completions": warmup_completions,
        "warmup_required": warmup_required,
        "warmup_gate_branch": warmup_gate_branch,
        "warmup_stability": warmup_stability,
        "drain_budget_seconds": drain_budget(pattern),
        "drain_quarantined": quarantined,
        "excluded_window_failures": failure_summary["excluded"],
        "excluded_window_failure_kinds": failure_summary["excluded_kinds"],
        "invariants": invariants,
    }


def validate_run(
    entry: dict[str, Any],
    profile: bool,
    burnin: bool,
    allow_missing_traces: bool = False,
) -> dict[str, Any]:
    run_dir = entry["run_dir"]
    client = validate_client(
        run_dir, entry["assignment"].cell.pattern, profile, burnin
    )
    layer1 = validate_layer1(run_dir)
    log = (run_dir / "server.log").read_text(errors="replace")
    details = CONFIGS[entry["assignment"].cell.config]
    server_invariants = {
        "triton_moe": "Using TRITON Unquantized MoE backend" in log,
        "chunked_mbt": (
            (
                details["mbt"] == 8192
                and "Chunked prefill is enabled with max_num_batched_tokens=8192"
                in log
            )
            or (
                details["mbt"] == 2048
                and "'max_num_batched_tokens': 2048" in log
                and "'enable_chunked_prefill': True" in log
            )
        ),
        "tp_effective": f"tensor_parallel_size={details['tp']}" in log,
        "drain_shutdown": "mode=drain timeout=120s" in log,
        "mns_effective": details["mns"] == 1024
        or "'max_num_seqs': 64" in log,
    }
    if not all(server_invariants.values()):
        raise RuntimeError(
            f"server config/backend failure: {run_dir}: {server_invariants}"
        )
    trace_summaries: list[dict[str, Any]] = []
    if profile:
        trace_dir = run_dir / "traces"
        traces = sorted(trace_dir.glob("*.pt.trace.json*"))
        expected = 2 * int(details["tp"])
        if len(traces) != expected and not allow_missing_traces:
            raise RuntimeError(
                f"trace count mismatch: {run_dir}: {len(traces)} != {expected}"
            )
        if traces:
            trace_summaries = [trace_summary(path) for path in traces]
            if not all(summary["valid"] for summary in trace_summaries):
                raise RuntimeError(f"invalid Layer-2 trace: {run_dir}")
            for rank in range(int(details["tp"])):
                if sum(f"rank{rank}" in path.name for path in traces) != 2:
                    raise RuntimeError(f"rank trace mismatch: {run_dir}: rank{rank}")
    forbidden = re.compile(r'"(?:prompt|messages|content|text)"\s*:')
    for path in (
        run_dir / "client/requests.jsonl",
        run_dir / "client/result.json",
        Path(layer1["stream"]),
    ):
        if forbidden.search(path.read_text(errors="replace")):
            raise RuntimeError(f"private/generated text field leaked: {path}")
    summary = {
        "schema": SCHEMA,
        "run_id": entry["run_id"],
        "pattern": entry["assignment"].cell.pattern,
        "config": entry["assignment"].cell.config,
        "gpus": entry["assignment"].gpus,
        "client": client,
        "layer1": layer1,
        "traces": trace_summaries,
        "missing_trace_files": (
            2 * int(details["tp"]) - len(trace_summaries) if profile else 0
        ),
        "layer2_missing_after_controller_cleanup": bool(
            profile and not trace_summaries and allow_missing_traces
        ),
        "drain_quarantined": client["drain_quarantined"],
        "server_invariants": server_invariants,
    }
    atomic_json(run_dir / "run-complete.json", summary)
    return summary


def run_stage(
    state: dict[str, Any],
    stage_name: str,
    assignments: list[Assignment],
    load_point: str,
    *,
    profile: bool,
    burnin: bool = False,
    confirmation: bool = False,
) -> None:
    stage_dir = RUN_ROOT / "stages" / stage_name
    existing = state["stages"].get(stage_name)
    if existing and existing.get("status") == "complete":
        if not (stage_dir / "stage-complete.json").exists():
            raise RuntimeError(f"complete stage marker missing: {stage_name}")
        print(f"RESUME skip {stage_name}", flush=True)
        return
    if stage_dir.exists():
        os.replace(stage_dir, stage_dir.with_name(f"{stage_dir.name}.interrupted-{int(time.time())}"))
    stage_dir.mkdir(parents=True)
    physical_gpus = sorted({gpu for item in assignments for gpu in item.gpus})
    reserve = sum(item.cell.width for item in assignments) * 1200 / 3600
    if state["gpu_hours_total"] + reserve > GPU_HOUR_LIMIT:
        raise RuntimeError(
            f"GPU budget reservation exceeds {GPU_HOUR_LIMIT}: "
            f"used={state['gpu_hours_total']:.6f} reserve={reserve:.6f}"
        )
    state["stages"][stage_name] = {
        "status": "preflight",
        "started_at": time.time(),
        "load_point": load_point,
        "profile": profile,
        "burnin": burnin,
        "confirmation": confirmation,
        "assignments": [
            {"cell": item.cell.cell_id, "gpus": item.gpus} for item in assignments
        ],
    }
    save_state(state)
    preflight(stage_dir, physical_gpus)
    entries: list[dict[str, Any]] = []
    handles: list[Any] = []
    owned_pgids: set[int] = set()
    monitor: common.Monitor | None = None
    failure: Exception | None = None
    stage_gpu_seconds = 0.0
    try:
        for assignment in assignments:
            suffix = "burnin" if burnin else ("confirmation" if confirmation else load_point)
            run_id = f"{assignment.cell.cell_id}-{suffix}"
            if burnin:
                run_dir = RUN_ROOT / "burnins" / assignment.cell.config
            elif confirmation:
                run_dir = RUN_ROOT / "confirmations" / assignment.cell.pattern
            else:
                run_dir = RUN_ROOT / "primary" / assignment.cell.cell_id / load_point
            if run_dir.exists():
                os.replace(
                    run_dir,
                    run_dir.with_name(f"{run_dir.name}.interrupted-{int(time.time())}"),
                )
            run_dir.mkdir(parents=True)
            trace_tmp = Path(f"/tmp/wjh-opprof-phase3-matrix/{run_id}")
            if trace_tmp.exists():
                shutil.rmtree(trace_tmp)
            trace_tmp.mkdir(parents=True)
            (run_dir / "trace-tmp").symlink_to(trace_tmp, target_is_directory=True)
            port = 8100 + assignment.gpus[0]
            command = server_command(assignment, port, trace_tmp)
            with (run_dir / "commands.log").open("w", encoding="utf-8") as output:
                output.write(
                    f"GPU_COMMAND {run_id} server: {shlex.join(command)}; "
                    "expected=startup 60-180s + fixed load 300-600s\n"
                )
            environment = os.environ.copy()
            environment.update(
                {
                    "CUDA_VISIBLE_DEVICES": ",".join(map(str, assignment.gpus)),
                    "VLLM_OPPROF_DIR": str(run_dir / "opprof"),
                    "HF_HUB_OFFLINE": "1",
                    "TRANSFORMERS_OFFLINE": "1",
                    "PYTHONUNBUFFERED": "1",
                }
            )
            handle = (run_dir / "server.log").open("ab", buffering=0)
            handles.append(handle)
            spawned_at = time.time()
            server = subprocess.Popen(
                command,
                cwd=SOURCE,
                env=environment,
                stdout=handle,
                stderr=subprocess.STDOUT,
                start_new_session=True,
            )
            owned_pgids.add(server.pid)
            entries.append(
                {
                    "assignment": assignment,
                    "run_id": run_id,
                    "run_dir": run_dir,
                    "trace_tmp": trace_tmp,
                    "port": port,
                    "server": server,
                    "spawned_at": spawned_at,
                }
            )
        state["stages"][stage_name]["status"] = "starting_servers"
        state["stages"][stage_name]["servers"] = {
            entry["run_id"]: {
                "pid": entry["server"].pid,
                "pgid": entry["server"].pid,
                "gpus": entry["assignment"].gpus,
            }
            for entry in entries
        }
        save_state(state)
        wait_ready(entries)
        monitor = common.Monitor(stage_dir / "monitor.jsonl", owned_pgids)
        monitor.start()
        for entry in entries:
            assignment = entry["assignment"]
            saturation_result = None
            if load_point == "moderate":
                saturation_result = (
                    RUN_ROOT
                    / "primary"
                    / assignment.cell.cell_id
                    / "saturation/client/result.json"
                )
            command = client_command(
                assignment,
                entry["port"],
                entry["run_dir"],
                load_point,
                profile,
                burnin,
                saturation_result,
            )
            with (entry["run_dir"] / "commands.log").open("a", encoding="utf-8") as output:
                output.write(
                    f"GPU_COMMAND {entry['run_id']} client: {shlex.join(command)}; "
                    "expected=fixed protocol timeline\n"
                )
            handle = (entry["run_dir"] / "client.log").open("ab", buffering=0)
            handles.append(handle)
            client = subprocess.Popen(
                command,
                cwd=WORKDIR,
                stdout=handle,
                stderr=subprocess.STDOUT,
                start_new_session=True,
            )
            entry["client"] = client
            owned_pgids.add(client.pid)
        state["stages"][stage_name]["status"] = "running_clients"
        state["stages"][stage_name]["clients"] = {
            entry["run_id"]: {"pid": entry["client"].pid, "pgid": entry["client"].pid}
            for entry in entries
        }
        save_state(state)
        deadline = time.monotonic() + 1200
        while time.monotonic() < deadline and any(
            entry["client"].poll() is None for entry in entries
        ):
            if any(entry["server"].poll() is not None for entry in entries):
                raise RuntimeError("server exited while a client was active")
            if monitor.other_apps:
                raise RuntimeError(f"other GPU process appeared: {monitor.other_apps}")
            live_hours = sum(
                (time.time() - entry["spawned_at"]) * entry["assignment"].cell.width
                for entry in entries
            ) / 3600
            if state["gpu_hours_total"] + live_hours >= GPU_HOUR_LIMIT:
                raise RuntimeError("cumulative GPU-hour hard stop reached")
            time.sleep(2)
        if any(entry["client"].poll() is None for entry in entries):
            raise TimeoutError(f"stage exceeded 1200 seconds: {stage_name}")
        bad = {}
        for entry in entries:
            if entry["client"].returncode == 0:
                continue
            sanity_path = entry["run_dir"] / "client/sanity.json"
            if not sanity_path.exists():
                bad[entry["run_id"]] = entry["client"].returncode
                continue
            client_sanity = json.loads(sanity_path.read_text())["invariants"]
            failed = [key for key, value in client_sanity.items() if not value]
            if failed != ["drain_within_timeout"]:
                bad[entry["run_id"]] = {
                    "returncode": entry["client"].returncode,
                    "failed_invariants": failed,
                }
        if bad:
            raise RuntimeError(f"client failures: {bad}")
        for entry in entries:
            if profile:
                destination = entry["run_dir"] / "traces"
                destination.mkdir()
                for path in sorted(entry["trace_tmp"].glob("*.pt.trace.json*")):
                    shutil.copy2(path, destination / path.name)
        state["stages"][stage_name]["status"] = "clients_complete"
        save_state(state)
    except Exception as error:
        failure = error
    finally:
        for entry in entries:
            client = entry.get("client")
            if client is not None and client.poll() is None:
                try:
                    os.killpg(client.pid, signal.SIGKILL)
                except ProcessLookupError:
                    pass
        try:
            stop_servers(entries)
        except Exception as error:
            failure = failure or error
        ended_at = time.time()
        stage_gpu_seconds = sum(
            (ended_at - entry["spawned_at"]) * entry["assignment"].cell.width
            for entry in entries
        )
        if monitor is not None:
            monitor.stop()
            atomic_json(stage_dir / "other-gpu-processes.json", monitor.other_apps)
            if monitor.other_apps:
                failure = failure or RuntimeError(
                    f"other GPU process appeared: {monitor.other_apps}"
                )
        for entry in entries:
            trace_link = entry["run_dir"] / "trace-tmp"
            if trace_link.is_symlink():
                trace_link.unlink()
            if entry["trace_tmp"].exists():
                shutil.rmtree(entry["trace_tmp"])
        for handle in handles:
            handle.close()
        (stage_dir / "clocks-after.txt").write_text(
            run_text(["nvidia-smi", "-q", "-d", "CLOCK"], check=False)
        )
        (stage_dir / "loadavg-after.txt").write_text(
            " ".join(str(value) for value in os.getloadavg()) + "\n"
        )
        try:
            common.verify_idle(physical_gpus, stage_dir)
        except Exception as error:
            failure = failure or error
    state["gpu_hours_total"] += stage_gpu_seconds / 3600
    state["gpu_hours_this_stage"] = stage_gpu_seconds / 3600
    if failure is None:
        try:
            summaries = [validate_run(entry, profile, burnin) for entry in entries]
        except Exception as error:
            failure = error
    if failure is not None:
        state["stages"][stage_name]["status"] = "failed"
        state["stages"][stage_name]["failure"] = repr(failure)
        state["stages"][stage_name]["gpu_hours"] = stage_gpu_seconds / 3600
        state["status"] = "failed"
        save_state(state)
        raise failure
    marker = {
        "schema": SCHEMA,
        "stage": stage_name,
        "completed_at": time.time(),
        "gpu_hours": stage_gpu_seconds / 3600,
        "runs": [summary["run_id"] for summary in summaries],
    }
    atomic_json(stage_dir / "stage-complete.json", marker)
    state["stages"][stage_name].update(
        {
            "status": "complete",
            "completed_at": marker["completed_at"],
            "gpu_hours": marker["gpu_hours"],
            "servers": {},
            "clients": {},
        }
    )
    if burnin:
        state["completed_burnins"] += len(assignments)
    else:
        state["completed_measured_runs"] += len(assignments)
        state.setdefault("drain_quarantined_runs", 0)
        state["drain_quarantined_runs"] += sum(
            summary["drain_quarantined"] for summary in summaries
        )
        if (
            state["completed_measured_runs"] > 0
            and state["drain_quarantined_runs"]
            / state["completed_measured_runs"]
            > 0.20
        ):
            state["status"] = "failed"
            save_state(state)
            raise RuntimeError("drain-quarantine rate exceeded 20%")
    save_state(state)


def cleanup_recorded(state: dict[str, Any]) -> None:
    for stage in state.get("stages", {}).values():
        if stage.get("status") == "complete":
            continue
        for group in ("clients", "servers"):
            for item in stage.get(group, {}).values():
                pgid = item.get("pgid")
                if pgid:
                    try:
                        os.killpg(int(pgid), signal.SIGKILL)
                    except ProcessLookupError:
                        pass
    time.sleep(2)


def execute(resume: bool) -> None:
    RUN_ROOT.mkdir(parents=True, exist_ok=True)
    state = load_state(resume)
    if resume:
        cleanup_recorded(state)
    current = fingerprint()
    if state["fingerprint"] and state["fingerprint"] != current:
        raise RuntimeError("resume fingerprint differs from frozen matrix")
    state["fingerprint"] = current
    state["status"] = "running"
    save_state(state)
    ensure_provenance()

    burnin_cells = [Cell("P06", config) for config in CONFIGS]
    for index, wave in enumerate(pack_cells(burnin_cells), 1):
        run_stage(
            state,
            f"burnin-{index:02d}",
            wave,
            "saturation",
            profile=False,
            burnin=True,
        )

    matrix_waves = pack_cells(cells())
    for index, wave in enumerate(matrix_waves, 1):
        run_stage(
            state,
            f"primary-{index:02d}-saturation",
            wave,
            "saturation",
            profile=True,
        )
        run_stage(
            state,
            f"primary-{index:02d}-moderate",
            wave,
            "moderate",
            profile=True,
        )

    confirmations = [
        Assignment(Cell(pattern, "C00"), (gpu,))
        for gpu, pattern in enumerate(("P10", "P06", "P03", "P01"))
    ]
    run_stage(
        state,
        "confirmations",
        confirmations,
        "moderate",
        profile=False,
        confirmation=True,
    )
    if state["completed_measured_runs"] != 52 or state["completed_burnins"] != 5:
        raise RuntimeError(
            f"completion count mismatch: measured={state['completed_measured_runs']} "
            f"burnins={state['completed_burnins']}"
        )
    trace_count = len(list((RUN_ROOT / "primary").glob("*/*/traces/*.pt.trace.json*")))
    expected_trace_count = 100 - state.get("missing_trace_files", 0)
    if trace_count != expected_trace_count:
        raise RuntimeError(
            f"trace aggregate mismatch: {trace_count} != {expected_trace_count}"
        )
    state["trace_files"] = trace_count
    state["status"] = "complete"
    state["completed_at"] = time.time()
    save_state(state)
    print(json.dumps(state, sort_keys=True), flush=True)


def plan() -> dict[str, Any]:
    matrix_waves = pack_cells(cells())
    return {
        "schema": SCHEMA,
        "placement": "4-way GPU0-3",
        "cells": len(cells()),
        "primary_runs": 48,
        "confirmation_runs": 4,
        "burnins": 5,
        "matrix_waves": [
            [
                {"cell": item.cell.cell_id, "gpus": item.gpus}
                for item in wave
            ]
            for wave in matrix_waves
        ],
        "expected_trace_files": 100,
        "prior_gpu_hours": PRIOR_GPU_HOURS,
        "gpu_hour_limit": GPU_HOUR_LIMIT,
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="command", required=True)
    run = subparsers.add_parser("run")
    run.add_argument("--resume", action="store_true")
    subparsers.add_parser("status")
    subparsers.add_parser("plan")
    args = parser.parse_args()
    if args.command == "run":
        execute(args.resume)
    elif args.command == "status":
        print(STATE.read_text() if STATE.exists() else '{"status":"absent"}')
    else:
        print(json.dumps(plan(), sort_keys=True, indent=2))


if __name__ == "__main__":
    main()