aituner/runs/opprof-phase6/opprof_phase6_controller.py

#!/usr/bin/env python3
"""Detached adaptive Phase-6 controller for the frozen 25-anchor surface."""

from __future__ import annotations

import argparse
import json
import os
import re
import shlex
import signal
import subprocess
import time
import urllib.request
from pathlib import Path
from typing import Any


WORKDIR = Path("/home/admin/cpfs/wjh/opprof-phase6-dash0-20260712")
RUN_ROOT = WORKDIR / "runs/phase6"
STATE = RUN_ROOT / "controller-state.json"
SOURCE = Path("/home/admin/cpfs/wjh/opprof-phase2-dash0-20260711/vllm-v0.24.0")
VENV = Path("/tmp/wjh-opprof-phase2-dash0-20260711/.venv")
AITUNER = Path("/home/admin/cpfs/wjh/aituner/aituner")
MODEL = Path("/home/admin/cpfs/wjh/models/Qwen/Qwen3-30B-A3B")
CLIENT = WORKDIR / "scripts/opprof_phase6_client.py"
PRIMARY_STUDY = WORKDIR / "provenance/study-primary.json"
TP4_STUDY = WORKDIR / "provenance/study-tp4.json"
GROUND = WORKDIR / "provenance/ground_truth.json"
GPU_LIMIT = 3.0
CPU_MAP = {i: f"{20*i}-{20*i+19}" for i in range(8)}
MARKER = "opprof-phase6-20260712"
OWNED_PGIDS: set[int] = set()


CELLS = {
    "tp1_mns8": {"tp": 1, "mns": 8, "lower": .21875, "peak": .2265625, "upper": .23046875},
    "tp1_mns16": {"tp": 1, "mns": 16, "lower": .2421875, "peak": .24609375, "upper": .25},
    "tp1_mns32": {"tp": 1, "mns": 32, "lower": .234375, "peak": .2421875, "upper": .24609375},
    "tp1_mns64": {"tp": 1, "mns": 64, "lower": .234375, "peak": .2421875, "upper": .24609375},
    "tp2_mns8": {"tp": 2, "mns": 8, "lower": .4921875, "peak": .49609375, "upper": .5},
    "tp2_mns16": {"tp": 2, "mns": 16, "lower": .4921875, "peak": .49609375, "upper": .5},
    "tp2_mns32": {"tp": 2, "mns": 32, "lower": .75, "peak": .75390625, "upper": .7578125},
    "tp2_mns64": {"tp": 2, "mns": 64, "lower": .5, "peak": .75, "upper": .75390625},
    "tp4_mns8": {"tp": 4, "mns": 8, "lower": .016055910008, "peak": .016591107009, "upper": .017126304009},
    "tp4_mns16": {"tp": 4, "mns": 16, "lower": .033182214016, "peak": .033717411016, "upper": .034252608017, "trap": True},
    "tp4_mns32": {"tp": 4, "mns": 32, "lower": .033182214016, "peak": .033717411016, "upper": .034252608017},
    "tp4_mns64": {"tp": 4, "mns": 64, "lower": .033182214016, "peak": .033717411016, "upper": .034252608017},
}
WAVES = [
    ("W1-tp1", [("tp1_mns8", (0,)), ("tp1_mns16", (1,)), ("tp1_mns32", (2,)), ("tp1_mns64", (3,))], .35),
    ("W2-tp2", [("tp2_mns8", (0,1)), ("tp2_mns16", (2,3)), ("tp2_mns32", (4,5)), ("tp2_mns64", (6,7))], .65),
    ("W3-tp4-trap", [("tp4_mns8", (0,1,2,3)), ("tp4_mns16", (4,5,6,7))], .85),
    ("W4-tp4", [("tp4_mns32", (0,1,2,3)), ("tp4_mns64", (4,5,6,7))], .65),
]


def atomic_json(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(json.dumps(value, sort_keys=True, indent=2) + "\n")
    os.replace(tmp, path)


def load_state() -> dict[str, Any]:
    if STATE.exists():
        return json.loads(STATE.read_text())
    return {
        "schema": 1, "status": "initialized", "gpu_hours_total": 0.0,
        "completed_primary_anchors": 0, "completed_confirmations": 0,
        "waves": {}, "failures": [], "started_at": time.time(),
    }


def save_state(state: dict[str, Any]) -> None:
    atomic_json(STATE, state)


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


def study_for(tp: int) -> Path:
    return TP4_STUDY if tp == 4 else PRIMARY_STUDY


def run_text(command: list[str], check: bool = True) -> str:
    result = subprocess.run(command, text=True, capture_output=True)
    if check and result.returncode:
        raise RuntimeError(f"command failed {command}: {result.stderr}")
    return result.stdout


def compute_pids() -> list[int]:
    text = run_text([
        "nvidia-smi", "--query-compute-apps=pid", "--format=csv,noheader,nounits"
    ], check=False)
    return sorted({int(x.strip()) for x in text.splitlines() if x.strip().isdigit()})


def pid_owned(pid: int) -> bool:
    try:
        if os.getpgid(pid) in OWNED_PGIDS:
            return True
    except ProcessLookupError:
        return True
    try:
        env = Path(f"/proc/{pid}/environ").read_bytes().split(b"\0")
    except (FileNotFoundError, PermissionError):
        return False
    return f"OPPROF_PHASE6_MARKER={MARKER}".encode() in env


def assert_no_other_compute() -> None:
    other = [pid for pid in compute_pids() if not pid_owned(pid)]
    if other:
        raise RuntimeError(f"outside GPU processes detected: {other}")


def assert_all_idle() -> None:
    if compute_pids():
        raise RuntimeError(f"GPU compute processes remain: {compute_pids()}")
    rows = run_text([
        "nvidia-smi", "--query-gpu=index,memory.used,utilization.gpu",
        "--format=csv,noheader,nounits",
    ])
    bad = []
    for line in rows.splitlines():
        index, memory, util = [int(x.strip()) for x in line.split(",")]
        if memory != 0 or util != 0:
            bad.append((index, memory, util))
    if bad:
        raise RuntimeError(f"GPU cleanup failure: {bad}")


def wait_ready(entry: dict[str, Any], timeout: float = 300.0) -> None:
    deadline = time.monotonic() + timeout
    url = f"http://127.0.0.1:{entry['port']}/v1/models"
    while time.monotonic() < deadline:
        if entry["server"].poll() is not None:
            raise RuntimeError(f"server exited before ready: {entry['cell']}")
        try:
            with urllib.request.urlopen(url, timeout=2) as response:
                if response.status < 500:
                    return
        except Exception:
            pass
        assert_no_other_compute()
        time.sleep(1)
    raise TimeoutError(f"server ready timeout: {entry['cell']}")


def server_command(cell: str, gpus: tuple[int, ...], port: int) -> list[str]:
    cfg = CELLS[cell]
    return [
        "taskset", "-c", cpu_mask(gpus), str(VENV / "bin/vllm"), "serve", str(MODEL),
        "--host", "127.0.0.1", "--port", str(port),
        "--served-model-name", "qwen3-30b-a3b-community",
        "--max-num-batched-tokens", "8192", "--max-num-seqs", str(cfg["mns"]),
        "--tensor-parallel-size", str(cfg["tp"]), "--shutdown-timeout", "120",
    ]


def client_command(entry: dict[str, Any], anchor: float, out: Path, warmup: bool) -> list[str]:
    cfg = CELLS[entry["cell"]]
    return [
        "taskset", "-c", cpu_mask(entry["gpus"]), str(VENV / "bin/python"), str(CLIENT),
        "warmup" if warmup else "run-anchor", "--study", str(study_for(cfg["tp"])),
        "--cell", entry["cell"], "--anchor", str(anchor), "--tp", str(cfg["tp"]),
        "--mns", str(cfg["mns"]), "--base-url", f"http://127.0.0.1:{entry['port']}",
        "--result-dir", str(out),
    ]


def live_gpu_hours(entries: list[dict[str, Any]]) -> float:
    now = time.time()
    return sum(
        len(e["gpus"]) * ((e.get("stopped_at") or now) - e["spawned_at"])
        for e in entries
    ) / 3600


def run_clients(
    entries: list[dict[str, Any]], assignments: list[tuple[dict[str, Any], float, Path]],
    state: dict[str, Any], wave_name: str, warmup: bool = False,
) -> list[dict[str, Any]]:
    processes = []
    handles = []
    for entry, anchor, out in assignments:
        command = client_command(entry, anchor, out, warmup)
        with (entry["dir"] / "commands.log").open("a") as f:
            f.write(f"CLIENT {shlex.join(command)}\n")
        handle = (out.parent / f"{out.name}.log").open("ab", buffering=0)
        handles.append(handle)
        client_env = os.environ.copy()
        client_env.update({"AITUNER_ROOT": str(AITUNER), "PYTHONUNBUFFERED": "1"})
        p = subprocess.Popen(
            command, cwd=WORKDIR, env=client_env, stdout=handle,
            stderr=subprocess.STDOUT, start_new_session=True,
        )
        processes.append((entry, anchor, out, p))
    deadline = time.monotonic() + 180
    while any(p.poll() is None for *_rest, p in processes):
        if time.monotonic() > deadline:
            raise TimeoutError(f"client batch timeout: {wave_name}")
        for entry in entries:
            if entry.get("stopped_at") is None and entry["server"].poll() is not None:
                raise RuntimeError(f"server exited during client: {entry['cell']}")
        assert_no_other_compute()
        if state["gpu_hours_total"] + live_gpu_hours(entries) >= GPU_LIMIT:
            raise RuntimeError("3.0 H20-hour hard stop reached")
        time.sleep(1)
    for handle in handles:
        handle.close()
    bad = [(e["cell"], a, p.returncode) for e, a, _o, p in processes if p.returncode]
    if bad:
        raise RuntimeError(f"client failures: {bad}")
    results = []
    for entry, anchor, out, _p in processes:
        result = json.loads((out / "result.json").read_text())
        results.append(result)
        entry.setdefault("results", []).append({"anchor": anchor, "dir": str(out), "kind": result["kind"]})
    return results


def stop_entry(entry: dict[str, Any]) -> None:
    if entry.get("stopped_at") is not None:
        return
    process = entry["server"]
    try:
        # Official vLLM shutdown: signal the API parent so EngineCore drains and
        # emits the in-stream footer/final sidecar. Process-group signals are
        # fallback only.
        os.kill(process.pid, signal.SIGINT)
    except ProcessLookupError:
        pass
    try:
        process.wait(timeout=150)
    except subprocess.TimeoutExpired:
        try:
            os.killpg(process.pid, signal.SIGTERM)
        except ProcessLookupError:
            pass
        try:
            process.wait(timeout=10)
        except subprocess.TimeoutExpired:
            try:
                os.killpg(process.pid, signal.SIGKILL)
            except ProcessLookupError:
                pass
            process.wait(timeout=30)
    entry["stopped_at"] = time.time()
    entry["server_handle"].close()


def validate_cell(entry: dict[str, Any]) -> dict[str, Any]:
    log = (entry["dir"] / "server.log").read_text(errors="replace").splitlines()
    ready = [i for i, line in enumerate(log) if "Application startup complete" in line]
    event = re.compile(r"torch\.compile took|Directly load AOT compilation|\bCompiling\b|Capturing CUDA graphs", re.I)
    post_ready_events = [i + 1 for i, line in enumerate(log) if event.search(line) and ready and i > ready[0]]
    streams = sorted((entry["dir"] / "opprof").glob("*.jsonl"))
    sidecars = sorted((entry["dir"] / "opprof").glob("*.jsonl.footer.json"))
    if len(streams) != 1 or len(sidecars) != 1:
        raise RuntimeError(f"Layer1 stream/sidecar mismatch: {entry['cell']}")
    decoded = [json.loads(line) for line in streams[0].read_text().splitlines()]
    footers = [item for item in decoded if item.get("record_type") == "footer"]
    records = [item for item in decoded if "step_index" in item]
    sidecar = json.loads(sidecars[0].read_text())
    indices = [int(item["step_index"]) for item in records]
    warm = json.loads((entry["dir"] / "warmup/result.json").read_text())
    intervals_ok = True
    for item in entry.get("results", []):
        if item["kind"] != "anchor":
            continue
        result = json.loads((Path(item["dir"]) / "result.json").read_text())
        lo, hi = result["interval"]["start_mono_ns"], result["interval"]["end_mono_ns"]
        intervals_ok &= any(lo <= int(r["submit_mono_ns"]) <= hi for r in records)
    common = {
        "one_ready_marker": len(ready) == 1,
        "compile_capture_pre_ready": not post_ready_events,
        "warmup_exact_16": warm["exact_output_count"] >= 16,
        "warmup_long": warm["selection"]["long_gt4096"] >= 1,
        "layer1_contiguous": indices == list(range(len(indices))),
        "written_matches_records": sidecar.get("written_records") == len(records),
        "encoded_balanced": sidecar.get("encoded_records") == sidecar.get("written_records") + sidecar.get("dropped_records"),
        "last_step_matches": bool(records) and sidecar.get("last_step_index") == records[-1]["step_index"],
        "layer1_zero_drops": sidecar.get("dropped_records") == 0,
        "anchor_intervals_present": intervals_ok,
    }
    if footers:
        accounting = {
            "one_footer_last": len(footers) == 1 and decoded[-1] is footers[0],
            "sidecar_final": sidecar.get("final") is True,
            "footer_sidecar_agrees": all(
                footers[0].get(key) == sidecar.get(key)
                for key in ("encoded_records", "written_records", "dropped_records")
            ),
        }
        accounting_mode = "graceful-footer"
    else:
        delta = abs(streams[0].stat().st_mtime_ns - int(sidecar["checkpoint_wall_ns"])) / 1e9
        accounting = {
            "checkpoint_sidecar": sidecar.get("final") is False,
            "checkpoint_within_flush_of_stream": delta <= float(sidecar["flush_interval_seconds"]) + .1,
        }
        accounting_mode = "checkpoint-sidecar-fallback"
    invariants = {**common, **accounting}
    if not all(invariants.values()):
        raise RuntimeError(f"cell validity failure {entry['cell']}: {invariants}")
    result = {
        "cell": entry["cell"], "invariants": invariants, "layer1_records": len(records),
        "stream": str(streams[0]), "post_ready_capture_events": post_ready_events,
        "accounting_mode": accounting_mode,
    }
    atomic_json(entry["dir"] / "cell-valid.json", result)
    return result


def historical_expected() -> dict[tuple[str, float], dict[str, Any]]:
    ground = json.loads(GROUND.read_text())
    result = {}
    for cell in ground["cells"]:
        for probe in cell["probe_history"]:
            result[(cell["cell_id"], float(probe["sampling_u"]))] = probe
    return result


def execute_wave(index: int, state: dict[str, Any], expected: dict[tuple[str, float], dict[str, Any]]) -> None:
    wave_name, assignments, estimate = WAVES[index]
    if state["waves"].get(wave_name, {}).get("status") == "complete":
        return
    future = sum(w[2] for w in WAVES[index:]) + .10
    if state["gpu_hours_total"] + future >= GPU_LIMIT:
        raise RuntimeError(f"projected budget exceeds cap before {wave_name}: {state['gpu_hours_total']+future}")
    echo = (
        f"WAVE_ECHO wave={wave_name} assignments="
        + ",".join(f"{cell}:gpu{'+'.join(map(str, gpus))}" for cell, gpus in assignments)
        + f" spent_h20h={state['gpu_hours_total']:.6f} wave_est_h20h={estimate:.3f} "
        + f"remaining_projection_h20h={future:.3f} cap_h20h={GPU_LIMIT:.1f} "
        + f"ground_truth={GROUND} workload=chat_w20260311_1000.jsonl"
    )
    with (RUN_ROOT / "launch-echo.log").open("a") as handle:
        handle.write(echo + "\n")
    print(echo, flush=True)
    assert_all_idle()
    wave_dir = RUN_ROOT / "waves" / wave_name
    wave_dir.mkdir(parents=True, exist_ok=True)
    entries = []
    state["status"] = "running"
    state["waves"][wave_name] = {"status": "starting", "estimate_h20_hours": estimate, "started_at": time.time()}
    save_state(state)
    failure = None
    try:
        for offset, (cell, gpus) in enumerate(assignments):
            cell_dir = RUN_ROOT / "cells" / cell
            cell_dir.mkdir(parents=True, exist_ok=True)
            port = 8500 + index * 10 + offset
            command = server_command(cell, gpus, port)
            with (cell_dir / "commands.log").open("a") as f:
                f.write(f"SERVER {shlex.join(command)}\n")
            handle = (cell_dir / "server.log").open("ab", buffering=0)
            env = os.environ.copy()
            env.update({
                "CUDA_VISIBLE_DEVICES": ",".join(map(str, gpus)),
                "VLLM_OPPROF_DIR": str(cell_dir / "opprof"),
                "OPPROF_PHASE6_MARKER": MARKER, "AITUNER_ROOT": str(AITUNER),
                "HF_HUB_OFFLINE": "1", "TRANSFORMERS_OFFLINE": "1", "PYTHONUNBUFFERED": "1",
            })
            server = subprocess.Popen(command, cwd=SOURCE, env=env, stdout=handle, stderr=subprocess.STDOUT, start_new_session=True)
            OWNED_PGIDS.add(server.pid)
            entries.append({
                "cell": cell, "gpus": gpus, "port": port, "dir": cell_dir,
                "server": server, "server_handle": handle, "spawned_at": time.time(),
            })
        for entry in entries:
            wait_ready(entry)
        state["waves"][wave_name]["status"] = "warmup"
        save_state(state)
        run_clients(entries, [
            (e, CELLS[e["cell"]]["peak"], e["dir"] / "warmup") for e in entries
        ], state, wave_name, warmup=True)
        state["waves"][wave_name]["status"] = "peaks"
        save_state(state)
        peaks = run_clients(entries, [
            (e, CELLS[e["cell"]]["peak"], e["dir"] / f"anchor-{CELLS[e['cell']]['peak']}")
            for e in entries
        ], state, wave_name)
        for result in peaks:
            old = expected[(result["cell"], float(result["anchor"]))]
            if result["selection"]["count"] != old["request_count"]:
                raise RuntimeError(f"selection mismatch {result['cell']} peak")
        neighbor_jobs = []
        for entry, result in zip(entries, peaks, strict=True):
            key = "upper" if result["feasible"] else "lower"
            anchor = CELLS[entry["cell"]][key]
            neighbor_jobs.append((entry, anchor, entry["dir"] / f"anchor-{anchor}"))
        neighbors = run_clients(entries, neighbor_jobs, state, wave_name)
        for result in neighbors:
            old = expected[(result["cell"], float(result["anchor"]))]
            if result["selection"]["count"] != old["request_count"]:
                raise RuntimeError(f"selection mismatch {result['cell']} neighbor")
        trap_entry = next((e for e in entries if CELLS[e["cell"]].get("trap")), None)
        if trap_entry is not None:
            used = {float(item["anchor"]) for item in trap_entry["results"] if item["kind"] == "anchor"}
            extra = next(CELLS[trap_entry["cell"]][k] for k in ("lower", "upper") if CELLS[trap_entry["cell"]][k] not in used)
            extra_result = run_clients(entries, [(trap_entry, extra, trap_entry["dir"] / f"anchor-{extra}")], state, wave_name)[0]
            if extra_result["selection"]["count"] != expected[(extra_result["cell"], float(extra))]["request_count"]:
                raise RuntimeError("trap extra selection mismatch")

        # Confirm only protocol-triggered anchors while the relevant server is hot.
        triggers = []
        for entry in entries:
            for item in entry.get("results", []):
                if item["kind"] != "anchor":
                    continue
                result = json.loads((Path(item["dir"]) / "result.json").read_text())
                old = expected[(entry["cell"], float(result["anchor"]))]
                flip = bool(result["feasible"]) != bool(old["feasible"])
                if .93 <= float(result["pass_rate"]) <= .97 or (entry["cell"] in {"tp2_mns32", "tp4_mns16"} and flip):
                    priority = 0 if entry["cell"] == "tp2_mns32" else (1 if entry["cell"] == "tp4_mns16" else 2)
                    triggers.append((priority, entry, float(result["anchor"])))
        triggers.sort(key=lambda x: x[0])
        for _priority, entry, anchor in triggers:
            projected_extra = len(entry["gpus"]) * 80 / 3600
            future_primary = sum(w[2] for w in WAVES[index + 1:])
            if state["gpu_hours_total"] + live_gpu_hours(entries) + future_primary + projected_extra + .03 >= GPU_LIMIT:
                state.setdefault("unconfirmed_triggers", []).append({"cell": entry["cell"], "anchor": anchor})
                continue
            confirm_index = 1 + sum(1 for item in entry.get("results", []) if item["kind"] == "anchor" and Path(item["dir"]).name.startswith("confirm"))
            out = entry["dir"] / f"confirm-{confirm_index}-anchor-{anchor}"
            run_clients(entries, [(entry, anchor, out)], state, wave_name)
            state["completed_confirmations"] += 1
        state["waves"][wave_name]["status"] = "stopping"
        save_state(state)
    except Exception as error:
        failure = error
    finally:
        for entry in entries:
            try:
                stop_entry(entry)
            except Exception as error:
                failure = failure or error
        time.sleep(2)
        try:
            assert_all_idle()
        except Exception as error:
            failure = failure or error
    wave_hours = live_gpu_hours(entries)
    state["gpu_hours_total"] += wave_hours
    state["waves"][wave_name]["gpu_hours"] = wave_hours
    if failure is not None:
        state["waves"][wave_name]["status"] = "failed"
        state["waves"][wave_name]["failure"] = repr(failure)
        state["status"] = "failed"
        state["failures"].append({"wave": wave_name, "failure": repr(failure)})
        save_state(state)
        raise failure
    validations = [validate_cell(entry) for entry in entries]
    primary_count = sum(
        1 for entry in entries for item in entry.get("results", [])
        if item["kind"] == "anchor" and not Path(item["dir"]).name.startswith("confirm")
    )
    state["completed_primary_anchors"] += primary_count
    state["waves"][wave_name].update({
        "status": "complete", "completed_at": time.time(), "primary_anchors": primary_count,
        "validations": validations,
    })
    save_state(state)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--resume", action="store_true")
    args = parser.parse_args()
    RUN_ROOT.mkdir(parents=True, exist_ok=True)
    state = load_state()
    expected = historical_expected()
    state["status"] = "running"
    save_state(state)
    for index in range(len(WAVES)):
        execute_wave(index, state, expected)
    state["status"] = "primary_complete"
    state["completed_at"] = time.time()
    save_state(state)
    print(json.dumps({
        "status": state["status"], "primary_anchors": state["completed_primary_anchors"],
        "confirmations": state["completed_confirmations"], "gpu_hours": state["gpu_hours_total"],
    }, sort_keys=True))


if __name__ == "__main__":
    main()