#!/usr/bin/env python3
"""Detached/resumable dash0 controller for the Phase-3 co-location gate."""

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 threading
import time
import urllib.request
from collections import defaultdict
from pathlib import Path
from typing import Any


SCHEMA = 1
WORKDIR = Path("/home/admin/cpfs/wjh/opprof-phase3-dash0-20260712")
RUN_ROOT = WORKDIR / "runs/e-a-colocation"
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"
CPU_MAP = {
    0: "0-19",
    1: "20-39",
    2: "40-59",
    3: "60-79",
    4: "80-99",
    5: "100-119",
    6: "120-139",
    7: "140-159",
}
BACKGROUND_8 = {
    1: "P01",
    2: "P02",
    3: "P03",
    4: "P04",
    5: "P06",
    6: "P07",
    7: "P08",
}
BACKGROUND_4 = {1: "P01", 2: "P03", 3: "P06"}
PATTERNS = {
    "P01": ["--kind", "synthetic", "--input-uniform", "128:512",
            "--output-fixed", "64", "--prefix", "none", "--arrival", "steady"],
    "P02": ["--kind", "synthetic", "--input-uniform", "128:512",
            "--output-fixed", "512", "--prefix", "none", "--arrival", "steady"],
    "P03": ["--kind", "synthetic", "--input-uniform", "4096:8192",
            "--output-fixed", "64", "--prefix", "none", "--arrival", "steady"],
    "P04": ["--kind", "synthetic", "--input-uniform", "4096:8192",
            "--output-fixed", "512", "--prefix", "none", "--arrival", "burst:8"],
    "P05": ["--kind", "synthetic", "--input-mixture",
            '{"uniform:128:512":0.5,"uniform:4096:8192":0.5}',
            "--output-fixed", "64", "--prefix", "none", "--arrival", "steady"],
    "P06": ["--kind", "synthetic", "--input-mixture",
            '{"uniform:128:512":0.5,"uniform:4096:8192":0.5}',
            "--output-fixed", "512", "--prefix", "none", "--arrival", "burst:8"],
    "P07": ["--kind", "synthetic", "--input-fixed", "1280",
            "--output-fixed", "512", "--prefix", "none", "--arrival", "burst:8"],
    "P08": ["--kind", "prefix-pool", "--num-prefixes", "8",
            "--prefix-len", "1024", "--suffix-fixed", "256",
            "--output-fixed", "512", "--arrival", "burst:8"],
}
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,
}


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def atomic_json(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_name(path.name + f".tmp.{os.getpid()}")
    with tmp.open("w", encoding="utf-8") as f:
        json.dump(value, f, sort_keys=True, indent=2)
        f.write("\n")
        f.flush()
        os.fsync(f.fileno())
    os.replace(tmp, path)


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"
            f"{result.stdout}"
        )
    return result.stdout


def command_log(path: Path, label: str, command: list[str], expected: str) -> None:
    with path.open("a", encoding="utf-8") as f:
        f.write(
            f"GPU_COMMAND {label}: {shlex.join(command)} ; expected={expected}\n"
        )


def load_state(resume: bool) -> dict[str, Any]:
    if STATE.exists():
        if not resume:
            raise RuntimeError("state exists; use --resume")
        return json.loads(STATE.read_text())
    return {
        "schema": SCHEMA,
        "status": "created",
        "created_at": time.time(),
        "controller_pid": os.getpid(),
        "stages": {},
        "fingerprint": {},
    }


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


def make_fingerprint() -> dict[str, Any]:
    return {
        "client_sha256": sha256_file(CLIENT),
        "controller_sha256": sha256_file(Path(__file__)),
        "source_commit": run_text(
            ["git", "-C", str(SOURCE), "rev-parse", "HEAD"]
        ).strip(),
        "model_path": str(MODEL),
        "venv_python": str(VENV / "bin/python"),
        # JSON object keys are strings. Normalize before the first write so a
        # read/compare during --resume is byte-semantically stable.
        "cpu_map": {str(gpu): cpus for gpu, cpus in CPU_MAP.items()},
    }


def ensure_provenance() -> None:
    provenance = RUN_ROOT / "provenance"
    provenance.mkdir(parents=True, exist_ok=True)
    for source in (CLIENT, Path(__file__)):
        dest = provenance / source.name
        source_sha = sha256_file(source)
        if dest.exists() and sha256_file(dest) != source_sha:
            # Preserve the prior immutable copy and add the reviewed runtime
            # repair as a content-addressed sibling.
            dest = provenance / f"{source.stem}.{source_sha[:12]}{source.suffix}"
        if dest.exists() and sha256_file(dest) != source_sha:
            raise RuntimeError(f"content-addressed provenance changed: {dest}")
        if not dest.exists():
            shutil.copy2(source, dest)
    checksums = {
        path.name: sha256_file(path)
        for path in sorted(provenance.glob("*.py"))
    }
    atomic_json(provenance / "sha256.json", checksums)


def ensure_manifests() -> None:
    PRIVATE.mkdir(parents=True, exist_ok=True, mode=0o700)
    os.chmod(PRIVATE, 0o700)
    for pattern, spec in PATTERNS.items():
        out = PRIVATE / f"{pattern}.jsonl"
        if out.exists():
            summary = json.loads(
                out.with_suffix(out.suffix + ".summary.json").read_text()
            )
            if summary["rows"] != 32768 or summary["sha256"] != sha256_file(out):
                raise RuntimeError(f"manifest verification failed: {out}")
            continue
        command = [
            str(VENV / "bin/python"),
            str(CLIENT),
            "materialize",
            "--id",
            pattern,
            *spec,
            "--num-requests",
            "32768",
            "--workload-seed",
            "20260712",
            "--out",
            str(out),
        ]
        run_text(command)


def gpu_query() -> list[dict[str, Any]]:
    text = run_text(
        [
            "nvidia-smi",
            "--query-gpu=index,uuid,memory.used,utilization.gpu",
            "--format=csv,noheader,nounits",
        ]
    )
    rows = []
    for line in text.strip().splitlines():
        index, uuid, memory, util = [part.strip() for part in line.split(",")]
        rows.append(
            {
                "index": int(index),
                "uuid": uuid,
                "memory_used_mib": int(memory),
                "utilization_pct": int(util),
            }
        )
    return rows


def compute_apps() -> list[dict[str, Any]]:
    text = run_text(
        [
            "nvidia-smi",
            "--query-compute-apps=gpu_uuid,pid,process_name,used_memory",
            "--format=csv,noheader,nounits",
        ],
        check=False,
    )
    rows = []
    for line in text.strip().splitlines():
        if not line or "No running" in line:
            continue
        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 preflight(gpus: list[int], stage_dir: Path) -> None:
    samples = []
    consecutive_idle = 0
    deadline = time.monotonic() + 60
    while time.monotonic() < deadline and consecutive_idle < 3:
        snapshots = [row for row in gpu_query() if row["index"] in gpus]
        apps = compute_apps()
        samples.append({"gpus": snapshots, "compute_apps": apps, "time": time.time()})
        idle = (
            not apps
            and all(row["memory_used_mib"] == 0 for row in snapshots)
            and all(row["utilization_pct"] == 0 for row in snapshots)
        )
        consecutive_idle = consecutive_idle + 1 if idle else 0
        if consecutive_idle < 3:
            time.sleep(2)
    atomic_json(stage_dir / "gpu-before.json", snapshots)
    atomic_json(stage_dir / "gpu-before-samples.json", samples)
    (stage_dir / "clocks-before.txt").write_text(
        run_text(["nvidia-smi", "-q", "-d", "CLOCK"])
    )
    (stage_dir / "loadavg-before.txt").write_text(
        " ".join(str(value) for value in os.getloadavg()) + "\n"
    )
    (stage_dir / "processes-before.txt").write_text(
        run_text(["ps", "-eo", "user,pid,ppid,pgid,lstart,args", "--sort=pid"])
    )
    if consecutive_idle < 3:
        raise RuntimeError(f"selected GPUs not stably idle: samples={samples}")


class Monitor:
    def __init__(self, path: Path, owned_pgids: set[int]) -> None:
        self.path = path
        self.owned_pgids = owned_pgids
        self.stop_event = threading.Event()
        self.other_apps: list[dict[str, Any]] = []
        self.thread = threading.Thread(target=self._run, daemon=True)

    def start(self) -> None:
        self.thread.start()

    def stop(self) -> None:
        self.stop_event.set()
        self.thread.join(timeout=20)

    def _run(self) -> None:
        with self.path.open("a", encoding="utf-8") as f:
            while not self.stop_event.is_set():
                apps = compute_apps()
                samples = []
                for app in apps:
                    try:
                        pgid = os.getpgid(app["pid"])
                    except ProcessLookupError:
                        pgid = None
                    sample = {**app, "pgid": pgid}
                    samples.append(sample)
                    if pgid not in self.owned_pgids:
                        self.other_apps.append(sample)
                clocks = run_text(
                    [
                        "nvidia-smi",
                        "--query-gpu=index,clocks.sm,clocks.mem,pstate,"
                        "utilization.gpu,memory.used",
                        "--format=csv,noheader,nounits",
                    ],
                    check=False,
                ).strip().splitlines()
                row = {
                    "wall_time": time.time(),
                    "loadavg": os.getloadavg(),
                    "gpu_clocks": clocks,
                    "compute_apps": samples,
                }
                f.write(json.dumps(row, sort_keys=True) + "\n")
                f.flush()
                self.stop_event.wait(15)


def server_command(gpu: int, stage_dir: Path, trace_dir: Path) -> list[str]:
    config = {**PROFILE_CONFIG, "torch_profiler_dir": str(trace_dir)}
    return [
        "taskset",
        "-c",
        CPU_MAP[gpu],
        str(VENV / "bin/vllm"),
        "serve",
        str(MODEL),
        "--host",
        "127.0.0.1",
        "--port",
        str(8000 + gpu),
        "--tensor-parallel-size",
        "1",
        "--enable-chunked-prefill",
        "--enable-prefix-caching",
        "--profiler-config",
        json.dumps(config, separators=(",", ":")),
    ]


def start_server(
    gpu: int, stage_dir: Path, state_stage: dict[str, Any]
) -> tuple[subprocess.Popen[Any], Any, Path]:
    gpu_dir = stage_dir / f"gpu{gpu}"
    gpu_dir.mkdir(parents=True, exist_ok=True)
    trace_dir = Path(f"/tmp/wjh-opprof-phase3-ea/{stage_dir.name}/gpu{gpu}")
    if trace_dir.exists():
        shutil.rmtree(trace_dir)
    trace_dir.mkdir(parents=True)
    command = server_command(gpu, stage_dir, trace_dir)
    command_log(
        stage_dir / "commands.log",
        f"{stage_dir.name} server gpu{gpu}",
        command,
        "startup 60-180s; run 300-450s",
    )
    env = os.environ.copy()
    env.update(
        {
            "CUDA_VISIBLE_DEVICES": str(gpu),
            "VLLM_OPPROF_DIR": str(gpu_dir / "opprof"),
            "HF_HUB_OFFLINE": "1",
            "TRANSFORMERS_OFFLINE": "1",
            "PYTHONUNBUFFERED": "1",
        }
    )
    log_handle = (gpu_dir / "server.log").open("ab", buffering=0)
    process = subprocess.Popen(
        command,
        cwd=SOURCE,
        env=env,
        stdout=log_handle,
        stderr=subprocess.STDOUT,
        start_new_session=True,
    )
    state_stage.setdefault("servers", {})[str(gpu)] = {
        "pid": process.pid,
        "pgid": process.pid,
        "command": command,
        "trace_dir": str(trace_dir),
    }
    return process, log_handle, trace_dir


def wait_ready(processes: dict[int, subprocess.Popen[Any]], timeout: float = 300) -> None:
    deadline = time.monotonic() + timeout
    pending = set(processes)
    while pending and time.monotonic() < deadline:
        for gpu in list(pending):
            if processes[gpu].poll() is not None:
                raise RuntimeError(f"server gpu{gpu} exited before ready")
            try:
                with urllib.request.urlopen(
                    f"http://127.0.0.1:{8000 + gpu}/health", timeout=1
                ) as response:
                    if response.status == 200:
                        pending.remove(gpu)
            except Exception:
                pass
        time.sleep(1)
    if pending:
        raise TimeoutError(f"servers not ready within {timeout}s: {sorted(pending)}")


def client_command(
    gpu: int,
    pattern: str,
    stage_dir: Path,
    load_point: str,
    saturation_result: Path | None,
    profile: bool,
    trace_dir: Path,
    post_clean_seconds: int,
) -> list[str]:
    result_dir = stage_dir / f"gpu{gpu}" / "client"
    command = [
        "taskset",
        "-c",
        CPU_MAP[gpu],
        str(VENV / "bin/python"),
        str(CLIENT),
        "run",
        "--manifest",
        str(PRIVATE / f"{pattern}.jsonl"),
        "--base-url",
        f"http://127.0.0.1:{8000 + gpu}",
        "--model",
        str(MODEL),
        "--load-point",
        load_point,
        "--max-concurrency",
        "256",
        "--ignore-eos",
        "--temperature",
        "0",
        "--warmup-seconds",
        "60",
        "--clean-segment-seconds",
        "80",
        "--num-clean-segments",
        "3",
        "--recovery-seconds",
        "30",
        "--drain-timeout-seconds",
        "120",
        "--workload-seed",
        "20260712",
        "--result-dir",
        str(result_dir),
    ]
    if load_point == "saturation":
        command += ["--request-rate", "inf"]
    else:
        assert saturation_result is not None
        command += [
            "--saturation-result",
            str(saturation_result),
            "--rate-fraction",
            "0.60",
        ]
    if profile:
        command += [
            "--profile-after-clean",
            "--num-profile-windows",
            "1",
            "--profile-warmup-iterations",
            "2",
            "--profile-active-iterations",
            "8",
            "--profile-trace-dir",
            str(trace_dir),
            "--profile-timeout-seconds",
            "120",
        ]
    elif post_clean_seconds:
        command += ["--post-clean-seconds", str(post_clean_seconds)]
    return command


def stop_processes(processes: list[subprocess.Popen[Any]]) -> None:
    live = [process for process in processes if process.poll() is None]
    for process in live:
        try:
            os.killpg(process.pid, signal.SIGINT)
        except ProcessLookupError:
            pass
    deadline = time.monotonic() + 60
    while time.monotonic() < deadline and any(p.poll() is None for p in live):
        time.sleep(1)
    for sig, wait_seconds in ((signal.SIGTERM, 20), (signal.SIGKILL, 5)):
        remaining = [p for p in live if p.poll() is None]
        if not remaining:
            break
        for process in remaining:
            try:
                os.killpg(process.pid, sig)
            except ProcessLookupError:
                pass
        deadline = time.monotonic() + wait_seconds
        while time.monotonic() < deadline and any(
            p.poll() is None for p in remaining
        ):
            time.sleep(0.5)


def _process_group_alive(pgid: int) -> bool:
    try:
        os.killpg(pgid, 0)
        return True
    except ProcessLookupError:
        return False


def stop_servers(processes: list[subprocess.Popen[Any]]) -> None:
    """Let the API parent ask EngineCore to close before group escalation."""
    groups = [process.pid for process in processes]
    for process in processes:
        if process.poll() is None:
            try:
                # Group SIGINT also interrupts EngineCore and loses the OpProf
                # footer. The API parent's shutdown path signals it cleanly.
                os.kill(process.pid, signal.SIGINT)
            except ProcessLookupError:
                pass
    deadline = time.monotonic() + 90
    while time.monotonic() < deadline and any(
        _process_group_alive(pgid) for pgid in groups
    ):
        time.sleep(1)
    for sig, wait_seconds in ((signal.SIGTERM, 20), (signal.SIGKILL, 5)):
        remaining = [pgid for pgid in groups if _process_group_alive(pgid)]
        if not remaining:
            break
        for pgid in remaining:
            try:
                os.killpg(pgid, sig)
            except ProcessLookupError:
                pass
        deadline = time.monotonic() + wait_seconds
        while time.monotonic() < deadline and any(
            _process_group_alive(pgid) for pgid in remaining
        ):
            time.sleep(0.5)


def verify_idle(gpus: list[int], stage_dir: Path) -> None:
    samples = []
    consecutive_zero = 0
    deadline = time.monotonic() + 120
    while time.monotonic() < deadline and consecutive_zero < 3:
        rows = [row for row in gpu_query() if row["index"] in gpus]
        apps = compute_apps()
        samples.append({"gpus": rows, "compute_apps": apps, "time": time.time()})
        zero = not apps and all(row["memory_used_mib"] == 0 for row in rows)
        consecutive_zero = consecutive_zero + 1 if zero else 0
        if consecutive_zero < 3:
            time.sleep(5)
    atomic_json(stage_dir / "gpu-after-cleanup.json", samples)
    if consecutive_zero < 3:
        raise RuntimeError("GPU memory did not reach three stable zero samples")


def run_stage(
    state: dict[str, Any],
    name: str,
    backgrounds: dict[int, str],
    load_point: str,
    saturation_stage: str | None,
    profile: bool,
) -> None:
    stage_dir = RUN_ROOT / name
    stage_record = state["stages"].get(name)
    if (
        stage_record
        and stage_record.get("status") == "complete"
        and (stage_dir / "stage-complete.json").exists()
    ):
        print(f"RESUME skip complete stage {name}", flush=True)
        return
    if stage_dir.exists():
        shutil.rmtree(stage_dir)
    stage_dir.mkdir(parents=True)
    gpus = [0, *sorted(backgrounds)]
    state["stages"][name] = {
        "status": "preflight",
        "started_at": time.time(),
        "gpus": gpus,
        "backgrounds": backgrounds,
        "load_point": load_point,
        "profile": profile,
    }
    save_state(state)
    preflight(gpus, stage_dir)
    servers: dict[int, subprocess.Popen[Any]] = {}
    server_logs = []
    client_processes: dict[int, subprocess.Popen[Any]] = {}
    client_logs = []
    trace_dirs: dict[int, Path] = {}
    monitor: Monitor | None = None
    failure: Exception | None = None
    try:
        for gpu in gpus:
            process, handle, trace_dir = start_server(
                gpu, stage_dir, state["stages"][name]
            )
            servers[gpu] = process
            server_logs.append(handle)
            trace_dirs[gpu] = trace_dir
        state["stages"][name]["status"] = "starting_servers"
        save_state(state)
        wait_ready(servers)
        state["stages"][name]["servers_ready_at"] = time.time()
        state["stages"][name]["status"] = "running_clients"
        save_state(state)
        monitor = Monitor(
            stage_dir / "monitor.jsonl", {process.pid for process in servers.values()}
        )
        monitor.start()

        assignments = {0: "P05", **backgrounds}
        for gpu, pattern in assignments.items():
            is_target = gpu == 0
            sat_result = (
                RUN_ROOT / saturation_stage / "gpu0/client/result.json"
                if is_target and saturation_stage
                else None
            )
            target_load = load_point if is_target else "saturation"
            command = client_command(
                gpu,
                pattern,
                stage_dir,
                target_load,
                sat_result,
                profile=is_target and profile,
                trace_dir=trace_dirs[gpu],
                post_clean_seconds=60 if not is_target and profile else 0,
            )
            command_log(
                stage_dir / "commands.log",
                f"{name} client gpu{gpu} {pattern}",
                command,
                "60s warmup + 240s clean + optional 2+8 profile/recovery",
            )
            handle = (stage_dir / f"gpu{gpu}/client.log").open("ab", buffering=0)
            client_logs.append(handle)
            process = subprocess.Popen(
                command,
                cwd=WORKDIR,
                stdout=handle,
                stderr=subprocess.STDOUT,
                start_new_session=True,
            )
            client_processes[gpu] = process
            state["stages"][name].setdefault("clients", {})[str(gpu)] = {
                "pid": process.pid,
                "pgid": process.pid,
                "command": command,
            }
        save_state(state)

        deadline = time.monotonic() + 780
        while time.monotonic() < deadline and any(
            process.poll() is None for process in client_processes.values()
        ):
            if any(process.poll() is not None for process in servers.values()):
                raise RuntimeError("server exited while client was active")
            time.sleep(2)
        if any(process.poll() is None for process in client_processes.values()):
            raise TimeoutError("client stage exceeded 780 seconds")
        bad = {
            gpu: process.returncode
            for gpu, process in client_processes.items()
            if process.returncode != 0
        }
        if bad:
            raise RuntimeError(f"client failures: {bad}")
        target_result = stage_dir / "gpu0/client/result.json"
        if not target_result.exists():
            raise RuntimeError("target result.json missing")
        if profile:
            trace_dest = stage_dir / "gpu0/traces"
            trace_dest.mkdir()
            for path in trace_dirs[0].glob("*"):
                if path.is_file():
                    shutil.copy2(path, trace_dest / path.name)
            traces = list(trace_dest.glob("*.pt.trace.json*"))
            if len(traces) != 1:
                raise RuntimeError(f"expected one target trace, found {len(traces)}")
        state["stages"][name]["status"] = "clients_complete"
        save_state(state)
    except Exception as exc:
        failure = exc
    finally:
        if monitor is not None:
            monitor.stop()
        stop_processes(list(client_processes.values()))
        stop_servers(list(servers.values()))
        for handle in client_logs + server_logs:
            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"
        )
        (stage_dir / "processes-after.txt").write_text(
            run_text(
                ["ps", "-eo", "user,pid,ppid,pgid,lstart,args", "--sort=pid"],
                check=False,
            )
        )
        if monitor is not None:
            atomic_json(stage_dir / "other-gpu-processes.json", monitor.other_apps)
        try:
            verify_idle(gpus, stage_dir)
        except Exception as exc:
            failure = failure or exc

    if failure is not None:
        state["stages"][name]["status"] = "failed"
        state["stages"][name]["failure"] = repr(failure)
        save_state(state)
        raise failure
    marker = {
        "schema": SCHEMA,
        "stage": name,
        "completed_at": time.time(),
        "target_result_sha256": sha256_file(
            stage_dir / "gpu0/client/result.json"
        ),
    }
    atomic_json(stage_dir / "stage-complete.json", marker)
    state["stages"][name].update(
        {"status": "complete", "completed_at": time.time(), "servers": {},
         "clients": {}}
    )
    save_state(state)


def classify_kernel(name: str) -> str:
    value = name.lower()
    if re.search(r"topkgating|moe_align|select_expert|moe.*rout", value):
        return "moe_router"
    if re.search(
        r"nccl|all_?reduce|reduce_scatter|all_?gather|custom.*all.*reduce", value
    ):
        return "collective"
    if re.search(
        r"flash|fmha|paged_attention|unified_attention|attention|"
        r"reshape_and_cache", value
    ):
        return "attention"
    if re.search(r"fused_moe|moe.*gemm|nvjet", value):
        return "moe_gemm"
    if re.search(
        r"sampling|(?<!moe_)topk|topp|argmax|multinomial|logits|penalty", value
    ):
        return "sampler"
    if re.search(r"gemm|matmul|cutlass", value):
        return "dense_gemm"
    if re.search(
        r"rms|layer.?norm|activation|residual|elementwise|reduce|fill|silu", value
    ):
        return "norm_elementwise"
    if re.search(r"cache|swap|memcpy|memset", value):
        return "kv_memory"
    return "other"


def trace_analysis(stage: str) -> dict[str, Any]:
    traces = list((RUN_ROOT / stage / "gpu0/traces").glob("*.pt.trace.json*"))
    if len(traces) != 1:
        raise RuntimeError(f"{stage}: expected one trace")
    trace = traces[0]
    opener = gzip.open if trace.suffix == ".gz" else open
    with opener(trace, "rt", encoding="utf-8") as f:
        data = json.load(f)
    durations: defaultdict[str, float] = defaultdict(float)
    kernel_count = 0
    steps = set()
    executes = 0
    for event in data.get("traceEvents", []):
        name = str(event.get("name", ""))
        match = re.fullmatch(r"ProfilerStep#(\d+)", name)
        if match:
            steps.add(int(match.group(1)))
        if name.startswith("execute_") and event.get("cat") == "user_annotation":
            executes += 1
        if event.get("cat") == "kernel" and float(event.get("dur", 0)) >= 0:
            kernel_count += 1
            durations[classify_kernel(name)] += float(event.get("dur", 0))
    total = sum(durations.values())
    shares = {
        family: duration / total for family, duration in sorted(durations.items())
    }
    result = {
        "schema": SCHEMA,
        "stage": stage,
        "trace_file": trace.name,
        "trace_sha256": sha256_file(trace),
        "trace_loadable": True,
        "kernel_count": kernel_count,
        "profiler_steps": sorted(steps),
        "execute_annotations": executes,
        "duration_us": dict(sorted(durations.items())),
        "shares": shares,
        "classifiable_fraction": 1.0 - shares.get("other", 0.0),
        "valid": (
            kernel_count > 0
            and steps == set(range(2, 10))
            and executes == 8
            and 1.0 - shares.get("other", 0.0) >= 0.70
        ),
    }
    return result


def compare_regime(prefix: str) -> dict[str, Any]:
    solo_result = json.loads(
        (RUN_ROOT / "solo-moderate/gpu0/client/result.json").read_text()
    )
    coloc_result = json.loads(
        (RUN_ROOT / f"{prefix}-moderate/gpu0/client/result.json").read_text()
    )
    solo_trace = trace_analysis("solo-moderate")
    coloc_trace = trace_analysis(f"{prefix}-moderate")
    solo_t = float(solo_result["clean"]["completed_throughput_rps"])
    coloc_t = float(coloc_result["clean"]["completed_throughput_rps"])
    top3 = sorted(
        solo_trace["shares"], key=solo_trace["shares"].get, reverse=True
    )[:3]
    share_deltas = {
        family: abs(
            coloc_trace["shares"].get(family, 0)
            - solo_trace["shares"].get(family, 0)
        )
        for family in top3
    }
    throughput_delta = abs(coloc_t / solo_t - 1)
    return {
        "schema": SCHEMA,
        "regime": prefix,
        "solo_throughput_rps": solo_t,
        "colocated_throughput_rps": coloc_t,
        "throughput_delta_fraction": throughput_delta,
        "top3_solo_families": top3,
        "solo_operator_shares": {
            family: solo_trace["shares"][family] for family in top3
        },
        "colocated_operator_shares": {
            family: coloc_trace["shares"].get(family, 0) for family in top3
        },
        "operator_share_delta_fraction": share_deltas,
        "solo_trace": solo_trace,
        "colocated_trace": coloc_trace,
        "pass": (
            solo_trace["valid"]
            and coloc_trace["valid"]
            and throughput_delta < 0.03
            and all(delta < 0.03 for delta in share_deltas.values())
        ),
    }


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


def write_analysis(comparisons: list[dict[str, Any]]) -> dict[str, Any]:
    accepted = next((c["regime"] for c in comparisons if c["pass"]), None)
    throughputs = [
        value
        for c in comparisons
        for value in (c["solo_throughput_rps"], c["colocated_throughput_rps"])
    ]
    deltas = [
        delta
        for c in comparisons
        for delta in c["operator_share_delta_fraction"].values()
    ]
    result = {
        "schema": SCHEMA,
        "comparisons": comparisons,
        "verdict": (
            f"{accepted.replace('coloc-', '')}-way"
            if accepted is not None
            else "STOP"
        ),
        "sanity": {
            "throughput_rps": numeric_sanity(throughputs),
            "throughput_delta": numeric_sanity(
                [c["throughput_delta_fraction"] for c in comparisons]
            ),
            "operator_share_delta": numeric_sanity(deltas),
            "invariants": {
                "throughputs_positive": all(x > 0 for x in throughputs),
                "share_deltas_in_0_1": all(0 <= x <= 1 for x in deltas),
                "top3_exact": all(
                    len(c["top3_solo_families"]) == 3 for c in comparisons
                ),
                "traces_valid": all(
                    c["solo_trace"]["valid"] and c["colocated_trace"]["valid"]
                    for c in comparisons
                ),
            },
        },
    }
    atomic_json(RUN_ROOT / "colocation-analysis.json", result)
    return result


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


def execute(resume: bool) -> None:
    RUN_ROOT.mkdir(parents=True, exist_ok=True)
    state = load_state(resume)
    if resume:
        clean_recorded_processes(state)
    fingerprint = make_fingerprint()
    if state["fingerprint"] and state["fingerprint"] != fingerprint:
        raise RuntimeError("resume fingerprint differs from frozen execution")
    state["fingerprint"] = fingerprint
    state["status"] = "running"
    save_state(state)
    ensure_provenance()
    ensure_manifests()
    run_stage(state, "solo-saturation", {}, "saturation", None, False)
    run_stage(
        state,
        "solo-moderate",
        {},
        "moderate",
        "solo-saturation",
        True,
    )
    run_stage(state, "coloc-8-saturation", BACKGROUND_8, "saturation", None, False)
    run_stage(
        state,
        "coloc-8-moderate",
        BACKGROUND_8,
        "moderate",
        "coloc-8-saturation",
        True,
    )
    comparisons = [compare_regime("coloc-8")]
    analysis = write_analysis(comparisons)
    if not analysis["comparisons"][0]["pass"]:
        run_stage(
            state, "coloc-4-saturation", BACKGROUND_4, "saturation", None, False
        )
        run_stage(
            state,
            "coloc-4-moderate",
            BACKGROUND_4,
            "moderate",
            "coloc-4-saturation",
            True,
        )
        comparisons.append(compare_regime("coloc-4"))
        analysis = write_analysis(comparisons)
    state["status"] = "complete" if analysis["verdict"] != "STOP" else "gate_failed"
    state["verdict"] = analysis["verdict"]
    state["completed_at"] = time.time()
    save_state(state)
    print(json.dumps(analysis, sort_keys=True), flush=True)


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("status")
    sub.add_parser("analyze")
    sub.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"}')
    elif args.command == "analyze":
        comparisons = [compare_regime("coloc-8")]
        if (RUN_ROOT / "coloc-4-moderate/stage-complete.json").exists():
            comparisons.append(compare_regime("coloc-4"))
        print(json.dumps(write_analysis(comparisons), sort_keys=True, indent=2))
    else:
        print(
            json.dumps(
                {
                    "target": "P05/C00 GPU0",
                    "background_8": BACKGROUND_8,
                    "background_4": BACKGROUND_4,
                    "cpu_map": CPU_MAP,
                    "stages": [
                        "solo-saturation",
                        "solo-moderate",
                        "coloc-8-saturation",
                        "coloc-8-moderate",
                        "conditional coloc-4 saturation/moderate",
                    ],
                },
                sort_keys=True,
                indent=2,
            )
        )


if __name__ == "__main__":
    main()