#!/usr/bin/env python3
"""Freeze Frontier's Qwen30 exact production-trace response surface."""

from __future__ import annotations

import argparse
import csv
import importlib.util
import json
import math
import os
import subprocess
import sys
import time
from pathlib import Path
from typing import Any


TARGET_PASS_RATE = 0.95
TPOT_SLOS_MS = (50.0, 100.0, 150.0, 180.0)
WINDOW_SECONDS = 600.0
GRAPH_CAPTURE_SIZES_BY_MNS = {
    8: (1, 2, 4, 8, 16),
    16: (1, 2, 4, 8, 16, 24, 32),
    32: (1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64),
    64: (1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128),
}
REAL_NUM_BLOCKS_BY_CONFIG = {
    (1, 8): 20137,
    (1, 16): 20128,
    (1, 32): 20108,
    (1, 64): 20069,
    (2, 8): 76639,
    (2, 16): 76620,
    (2, 32): 76583,
    (2, 64): 76505,
    (4, 8): 191930,
    (4, 16): 191882,
    (4, 32): 191786,
    (4, 64): 191589,
}
KERNEL_DECODE_KV_CONTEXTS = (128, 1024, 2048, 4096, 8192, 16384, 32768, 40960)
BASE_RUNNER = (
    Path(__file__).resolve().parents[1]
    / "frontier-phase-factorial-v0/run_frontier_qwen30_prefill_surface.py"
)


def load_base():
    spec = importlib.util.spec_from_file_location("qwen30_prefill_surface_base", BASE_RUNNER)
    if spec is None or spec.loader is None:
        raise ImportError(BASE_RUNNER)
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


BASE = load_base()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--frontier-source", type=Path, required=True)
    parser.add_argument("--replayserve-root", type=Path, required=True)
    parser.add_argument("--profile-root", type=Path, required=True)
    parser.add_argument("--kernel-profile-root", type=Path)
    parser.add_argument("--python-deps", type=Path, required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    parser.add_argument(
        "--predictor-cache-root",
        type=Path,
        help="Reuse a validated predictor cache outside the output root.",
    )
    parser.add_argument(
        "--trace",
        action="append",
        required=True,
        help="Frozen trace anchor as LABEL=PATH; repeat in increasing load order.",
    )
    parser.add_argument("--config", action="append")
    parser.add_argument(
        "--rate-contract",
        choices=("trace-window", "uniform-spacing"),
        default="trace-window",
    )
    parser.add_argument(
        "--prefix-caching",
        action=argparse.BooleanOptionalAction,
        default=True,
    )
    parser.add_argument(
        "--cc-backend", choices=("analytical", "vidur"), default="vidur"
    )
    parser.add_argument("--allreduce-csv", type=Path)
    parser.add_argument("--timeout-seconds", type=float, default=1800.0)
    parser.add_argument("--predictor-training-job-threads", type=int, default=1)
    parser.add_argument(
        "--decode-cuda-graph-mode",
        choices=("none", "full_decode_only", "piecewise"),
        default="none",
    )
    parser.add_argument(
        "--align-real-graph-runtime",
        action="store_true",
        help="Use real observed capture lists and per-(TP,MNS) KV blocks.",
    )
    parser.add_argument(
        "--fresh-predictor-cache",
        action="store_true",
        help="Disable Frontier predictor cache reuse for this profile family.",
    )
    parser.add_argument("--resume", action="store_true")
    parser.add_argument("--continue-on-failure", action="store_true")
    parser.add_argument(
        "--store-stage-batch-ledger",
        action="store_true",
        help="Persist Frontier's per-stage batch/state ledger for diagnosis.",
    )
    return parser.parse_args()


def ttft_limit_ms(input_tokens: int) -> float:
    return 1000.0 + 1000.0 * input_tokens / 8000.0


def percentile(values: list[float], fraction: float) -> float | None:
    if not values:
        return None
    ordered = sorted(values)
    return ordered[math.ceil(fraction * len(ordered)) - 1]


def validate_allreduce_profile(
    path: Path | None, required_tps: set[int]
) -> dict[str, Any] | None:
    """Fail before simulation when a measured collective profile misses a TP."""
    if not required_tps and path is None:
        return None
    if path is None:
        raise ValueError(
            f"measured all-reduce profile is required for TP coverage {sorted(required_tps)}"
        )
    if not path.is_file():
        raise FileNotFoundError(path)

    required_fields = {
        "time_stats.all_reduce.median",
        "num_workers",
        "size",
        "collective",
    }
    rows_by_tp: dict[int, int] = {}
    sizes_by_tp: dict[int, set[int]] = {}
    seen: set[tuple[int, int]] = set()
    with path.open(newline="") as handle:
        reader = csv.DictReader(handle)
        missing_fields = required_fields - set(reader.fieldnames or ())
        if missing_fields:
            raise ValueError(
                f"all-reduce profile schema missing fields: {sorted(missing_fields)}"
            )
        for line_number, row in enumerate(reader, start=2):
            if row["collective"] != "all_reduce":
                raise ValueError(
                    f"unexpected collective at {path}:{line_number}: {row['collective']!r}"
                )
            tp = int(row["num_workers"])
            size = int(row["size"])
            median = float(row["time_stats.all_reduce.median"])
            if size <= 0 or not math.isfinite(median) or median <= 0:
                raise ValueError(
                    f"invalid measured all-reduce row at {path}:{line_number}"
                )
            key = (tp, size)
            if key in seen:
                raise ValueError(f"duplicate measured all-reduce row: TP{tp}, size={size}")
            seen.add(key)
            rows_by_tp[tp] = rows_by_tp.get(tp, 0) + 1
            sizes_by_tp.setdefault(tp, set()).add(size)

    missing_tps = required_tps - set(rows_by_tp)
    if missing_tps:
        raise ValueError(
            "measured all-reduce profile coverage missing "
            f"TPs {sorted(missing_tps)}; available={sorted(rows_by_tp)}; "
            "analytical fallback is forbidden"
        )
    return {
        "required_tp_coverage": sorted(required_tps),
        "available_tp_coverage": sorted(rows_by_tp),
        "rows_by_tp": {str(tp): rows_by_tp[tp] for tp in sorted(rows_by_tp)},
        "size_range_by_tp": {
            str(tp): [min(sizes), max(sizes)]
            for tp, sizes in sorted(sizes_by_tp.items())
        },
    }


def collective_fallback_evidence(run_dir: Path) -> list[dict[str, Any]]:
    patterns = (
        "all-reduce profiling data is empty after filtering",
        "analytical fallback",
        "fallback to analytical",
    )
    evidence = []
    for name in ("stdout.log", "stderr.log"):
        path = run_dir / name
        for line_number, line in enumerate(
            path.read_text(errors="replace").splitlines(), start=1
        ):
            if any(pattern in line.lower() for pattern in patterns):
                evidence.append(
                    {"log": name, "line_number": line_number, "text": line.strip()}
                )
    return evidence


def manifest_offered_rate(path: Path) -> tuple[float | None, str | None]:
    manifest_path = path.with_name("manifest.json")
    if not manifest_path.is_file():
        return None, None
    manifest = json.loads(manifest_path.read_text())
    if manifest.get("public_csv") != str(path):
        return None, None
    rate = manifest.get("global_offered_request_rate")
    if not isinstance(rate, (int, float)) or not math.isfinite(rate) or rate <= 0:
        raise ValueError(f"invalid global_offered_request_rate in {manifest_path}")
    return float(rate), str(manifest_path)


def parse_trace(
    specification: str, *, rate_contract: str = "trace-window", prefix_caching: bool
) -> dict[str, Any]:
    if "=" not in specification:
        raise ValueError(f"trace must be LABEL=PATH: {specification}")
    label, raw_path = specification.split("=", 1)
    if not label or "/" in label:
        raise ValueError(f"invalid trace label: {label!r}")
    path = Path(raw_path).resolve()
    if not path.is_file():
        raise FileNotFoundError(path)
    with path.open(newline="") as source:
        rows = list(csv.DictReader(source))
    if not rows:
        raise ValueError(f"empty trace: {path}")
    required = {
        "arrived_at",
        "num_prefill_tokens",
        "num_decode_tokens",
        "session_id",
        "block_hash_ids",
    }
    if not required.issubset(rows[0]):
        raise ValueError(f"trace columns missing: {required - set(rows[0])}")
    arrivals = [float(row["arrived_at"]) for row in rows]
    if any(not math.isfinite(value) or value < 0 for value in arrivals):
        raise ValueError(f"invalid arrival in {path}")
    if any(right < left for left, right in zip(arrivals, arrivals[1:])):
        raise ValueError(f"arrival order drift in {path}")
    if rate_contract == "trace-window":
        offered_request_rate, rate_manifest = manifest_offered_rate(path)
        if offered_request_rate is None:
            offered_request_rate = len(rows) / WINDOW_SECONDS
            rate_source = "legacy_fixed_600_second_window"
        else:
            rate_source = f"manifest:{rate_manifest}"
    elif rate_contract == "uniform-spacing":
        if len(rows) < 2 or arrivals[-1] <= arrivals[0]:
            raise ValueError("uniform-spacing traces require at least two arrivals")
        intervals = [right - left for left, right in zip(arrivals, arrivals[1:])]
        expected_interval = (arrivals[-1] - arrivals[0]) / (len(arrivals) - 1)
        if any(abs(value - expected_interval) > 1e-9 for value in intervals):
            raise ValueError(f"non-uniform fixed trace: {path}")
        offered_request_rate = 1.0 / expected_interval
        rate_source = "uniform_spacing"
    else:
        raise ValueError(f"unknown rate contract: {rate_contract}")
    shapes = [
        (int(row["num_prefill_tokens"]), int(row["num_decode_tokens"]))
        for row in rows
    ]
    if any(prefill <= 0 or decode <= 0 or prefill + decode > 40960 for prefill, decode in shapes):
        raise ValueError(f"out-of-contract shape in {path}")
    if prefix_caching:
        for index, (row, (prefill, _)) in enumerate(zip(rows, shapes, strict=True)):
            block_ids = [value for value in row["block_hash_ids"].split("|") if value]
            if len(block_ids) != prefill // 16:
                raise ValueError(
                    "Frontier prefix-cache trace must expose only complete "
                    f"16-token blocks: row={index}, ISL={prefill}, "
                    f"ids={len(block_ids)}, expected={prefill // 16}"
                )
    return {
        "label": label,
        "path": path,
        "sha256": BASE.sha256(path),
        "requests": len(rows),
        "offered_request_rate": offered_request_rate,
        "offered_request_rate_source": rate_source,
        "first_arrival_s": arrivals[0],
        "last_arrival_s": arrivals[-1],
        "shapes": shapes,
    }


def find_request_metrics(run_dir: Path) -> Path:
    matches = list((run_dir / "metrics").rglob("request_metrics.csv"))
    if len(matches) != 1:
        raise RuntimeError(f"expected one request_metrics.csv, got {matches}")
    return matches[0]


def trace_manifest_entry(trace: dict[str, Any]) -> dict[str, Any]:
    return {
        key: str(value) if isinstance(value, Path) else value
        for key, value in trace.items()
        if key != "shapes"
    }


def classify_frontier_failure(stderr: str) -> str:
    if "Sequential simulation ended with non-empty scheduler state" in stderr:
        return "scheduler_stall"
    return "frontier_error"


def configure_stage_batch_ledger(command: list[str], *, enabled: bool) -> list[str]:
    if not enabled:
        return command
    disabled = "--no-metrics_config_store_frontier_stage_batch_ledger"
    enabled_flag = "--metrics_config_store_frontier_stage_batch_ledger"
    if command.count(disabled) != 1 or enabled_flag in command:
        raise ValueError("unexpected stage-batch-ledger command contract")
    command[command.index(disabled)] = enabled_flag
    return command


def score(path: Path, expected_shapes: list[tuple[int, int]]) -> dict[str, Any]:
    with path.open(newline="") as source:
        rows = list(csv.DictReader(source))
    if len(rows) != len(expected_shapes):
        raise ValueError(f"request count mismatch: {len(rows)} != {len(expected_shapes)}")
    rows.sort(key=lambda row: int(row["Request Id"]))
    request_metrics = []
    for index, (row, expected) in enumerate(zip(rows, expected_shapes, strict=True)):
        if int(row["Request Id"]) != index:
            raise ValueError("Frontier request ID/order drift")
        shape = (
            int(float(row["request_num_prefill_tokens"])),
            int(float(row["request_num_decode_tokens"])),
        )
        if shape != expected:
            raise ValueError(f"request shape drift at {index}: {shape} != {expected}")
        ttft = float(row["ttft"])
        e2e = float(row["request_e2e_time"])
        tpot = (e2e - ttft) / (shape[1] - 1) if shape[1] > 1 else None
        values = [ttft, e2e] + ([] if tpot is None else [tpot])
        if not all(math.isfinite(value) and value >= 0 for value in values):
            raise ValueError(f"invalid latency at request {index}")
        request_metrics.append(
            {
                "request_id": index,
                "input_tokens": shape[0],
                "output_tokens": shape[1],
                "ttft_ms": ttft,
                "ttft_limit_ms": ttft_limit_ms(shape[0]),
                "tpot_ms": tpot,
                "e2e_ms": e2e,
            }
        )
    slos = {}
    for limit in TPOT_SLOS_MS:
        passed = sum(
            row["ttft_ms"] <= row["ttft_limit_ms"]
            and (row["tpot_ms"] is None or row["tpot_ms"] <= limit)
            for row in request_metrics
        )
        pass_rate = passed / len(request_metrics)
        slos[f"tpot_{int(limit)}ms"] = {
            "passed": passed,
            "pass_rate": pass_rate,
            "feasible": pass_rate >= TARGET_PASS_RATE,
        }
    ttfts = [float(row["ttft_ms"]) for row in request_metrics]
    tpots = [float(row["tpot_ms"]) for row in request_metrics if row["tpot_ms"] is not None]
    e2es = [float(row["e2e_ms"]) for row in request_metrics]
    def summary(values: list[float]) -> dict[str, float | None]:
        if not values:
            return {"mean": None, "p50": None, "p90": None, "p95": None}
        return {
            "mean": sum(values) / len(values),
            "p50": percentile(values, 0.50),
            "p90": percentile(values, 0.90),
            "p95": percentile(values, 0.95),
        }

    ttft_summary = summary(ttfts)
    tpot_summary = summary(tpots)
    e2e_summary = summary(e2es)
    return {
        **{f"ttft_{name}_ms": value for name, value in ttft_summary.items()},
        **{f"tpot_{name}_ms": value for name, value in tpot_summary.items()},
        **{f"e2e_{name}_ms": value for name, value in e2e_summary.items()},
        "slos": slos,
    }


def kernel_profile_paths(root: Path) -> dict[str, Path]:
    paths = {
        "linear": root / "linear_op.csv",
        "attention": root / "attention.csv",
        "moe": root / "moe.csv",
        "manifest": root / "manifest.json",
    }
    missing = [str(path) for path in paths.values() if not path.is_file()]
    if missing:
        raise FileNotFoundError(missing)
    return paths


def validate_kernel_profile(paths: dict[str, Path]) -> dict[str, Any]:
    manifest = json.loads(paths["manifest"].read_text())
    outputs = manifest.get("outputs", {})
    for filename, name in (
        ("linear_op.csv", "linear"),
        ("attention.csv", "attention"),
        ("moe.csv", "moe"),
    ):
        if outputs.get(filename) != BASE.sha256(paths[name]):
            raise ValueError(f"kernel-only profile hash mismatch for {filename}")

    with paths["linear"].open(newline="") as source:
        linear_rows = list(csv.DictReader(source))
    with paths["attention"].open(newline="") as source:
        attention_rows = list(csv.DictReader(source))
    with paths["moe"].open(newline="") as source:
        moe_rows = list(csv.DictReader(source))
    for label, rows in (("linear", linear_rows), ("attention", attention_rows), ("moe", moe_rows)):
        if not rows or {row.get("measurement_type") for row in rows} != {"KERNEL_ONLY"}:
            raise ValueError(f"{label} lacks an exclusive KERNEL_ONLY measurement family")

    required_buckets = set(GRAPH_CAPTURE_SIZES_BY_MNS[64])
    coverage: dict[str, Any] = {}
    for tp in (1, 2, 4):
        linear_tokens = {
            int(float(row["num_tokens"]))
            for row in linear_rows
            if int(float(row["num_tensor_parallel_workers"])) == tp
        }
        moe_tokens = {
            int(float(row["num_tokens"]))
            for row in moe_rows
            if int(float(row["num_tensor_parallel_workers"])) == tp
        }
        attention_pairs = {
            (int(float(row["batch_size"])), int(float(row["kv_cache_size"])))
            for row in attention_rows
            if int(float(row["num_tensor_parallel_workers"])) == tp
            and row["is_prefill"].lower() == "false"
            and row.get("is_true_mixed_batch", "").lower() != "true"
        }
        missing_linear = required_buckets - linear_tokens
        missing_moe = required_buckets - moe_tokens
        missing_attention = {
            (bucket, kv)
            for bucket in required_buckets
            for kv in KERNEL_DECODE_KV_CONTEXTS
            if (bucket, kv) not in attention_pairs
        }
        if missing_linear or missing_moe or missing_attention:
            raise ValueError(
                f"kernel-only profile coverage TP{tp}: linear={sorted(missing_linear)}, "
                f"moe={sorted(missing_moe)}, attention={sorted(missing_attention)}"
            )
        coverage[str(tp)] = {
            "linear_tokens": sorted(linear_tokens),
            "moe_tokens": sorted(moe_tokens),
            "attention_decode_pairs": len(attention_pairs),
        }
    return {"manifest": manifest, "coverage": coverage}


def main() -> None:
    args = parse_args()
    if args.predictor_training_job_threads <= 0:
        raise ValueError("predictor training job threads must be positive")
    for name in (
        "frontier_source",
        "replayserve_root",
        "profile_root",
        "python_deps",
        "output_root",
    ):
        setattr(args, name, getattr(args, name).resolve())
    if args.allreduce_csv is not None:
        args.allreduce_csv = args.allreduce_csv.resolve()
    if args.kernel_profile_root is not None:
        args.kernel_profile_root = args.kernel_profile_root.resolve()
        if args.decode_cuda_graph_mode == "none":
            raise ValueError("--kernel-profile-root requires a non-none graph mode")
    traces = [
        parse_trace(
            specification,
            rate_contract=args.rate_contract,
            prefix_caching=args.prefix_caching,
        )
        for specification in args.trace
    ]
    if len({trace["label"] for trace in traces}) != len(traces):
        raise ValueError("trace labels must be unique")
    if any(
        right["offered_request_rate"] <= left["offered_request_rate"]
        for left, right in zip(traces, traces[1:])
    ):
        raise ValueError("trace anchors must be supplied in increasing load order")

    selected = list(BASE.GRID)
    if args.config:
        wanted = set(args.config)
        selected = [config for config in BASE.GRID if config.name in wanted]
        if {config.name for config in selected} != wanted:
            raise ValueError(f"unknown configs: {wanted - {config.name for config in selected}}")
    allreduce_profile = None
    if args.cc_backend == "vidur":
        allreduce_profile = validate_allreduce_profile(
            args.allreduce_csv, {config.tp for config in selected if config.tp > 1}
        )
    allreduce_csv_sha256 = (
        BASE.sha256(args.allreduce_csv) if args.allreduce_csv else None
    )
    paths = BASE.profile_paths(args.profile_root)
    coverage = BASE.validate_profile(paths)
    kernel_paths = None
    kernel_coverage = None
    if args.kernel_profile_root is not None:
        kernel_paths = kernel_profile_paths(args.kernel_profile_root)
        kernel_coverage = validate_kernel_profile(kernel_paths)
    builder = BASE.load_module(
        "qwen30_exact_trace_frontier_builder",
        args.replayserve_root / "tools/run_frontier_sweep.py",
    )
    frontier_head = subprocess.run(
        ["git", "-C", str(args.frontier_source), "rev-parse", "HEAD"],
        check=True,
        text=True,
        stdout=subprocess.PIPE,
    ).stdout.strip()
    environment = os.environ.copy()
    pythonpath = [str(args.python_deps), str(args.frontier_source)]
    if environment.get("PYTHONPATH"):
        pythonpath.append(environment["PYTHONPATH"])
    environment.update(
        {
            "PYTHONPATH": ":".join(pythonpath),
            "CUDA_VISIBLE_DEVICES": "",
            "NVIDIA_VISIBLE_DEVICES": "void",
            "WANDB_DISABLED": "true",
            "VIDUR_DISABLE_WANDB": "1",
            "FRONTIER_LOG_LEVEL": "WARNING",
            "PYTHONDONTWRITEBYTECODE": "1",
        }
    )

    config_results = []
    predictor_cache_root = args.predictor_cache_root or args.output_root / "cache"
    for config in selected:
        loads = []
        config_knobs = BASE.knobs(config, paths, predictor_cache_root)
        config_knobs["enable_prefix_caching"] = args.prefix_caching
        config_knobs["prediction_max_tokens_per_request"] = 40960
        config_knobs["decode_cuda_graph_mode"] = args.decode_cuda_graph_mode
        config_knobs["no_cache"] = args.fresh_predictor_cache
        if args.align_real_graph_runtime:
            config_knobs["num_blocks"] = REAL_NUM_BLOCKS_BY_CONFIG[(config.tp, config.mns)]
            # These are predictor-grid limits, not scheduler knobs.  Keep the
            # grid inside the observed vLLM graph/runtime envelope: no batch
            # can use a capture bucket above this cell's list and MBT is 8192.
            config_knobs["prediction_max_batch_size"] = max(
                GRAPH_CAPTURE_SIZES_BY_MNS[config.mns]
            )
            config_knobs["prediction_max_prefill_chunk_size"] = 8192
        if kernel_paths is not None:
            config_knobs.update(
                {
                    "linear_op_kernel_only_input_file": str(kernel_paths["linear"]),
                    "atten_kernel_only_input_file": str(kernel_paths["attention"]),
                    "moe_kernel_only_input_file": str(kernel_paths["moe"]),
                }
            )
        for trace in traces:
            run_dir = args.output_root / "runs" / config.name / trace["label"]
            result_path = run_dir / "result.json"
            if args.resume and result_path.is_file():
                previous = json.loads(result_path.read_text())
                if (
                    previous.get("status") == "completed"
                    and previous.get("allreduce_csv_sha256")
                    == allreduce_csv_sha256
                ):
                    loads.append(previous)
                    continue
            run_dir.mkdir(parents=True, exist_ok=True)
            command = builder.build_frontier_command(
                python_bin="/usr/bin/python3",
                trace_file=trace["path"],
                metrics_root=run_dir / "metrics",
                run_id=f"qwen30_trace_{config.name}_{trace['label']}",
                knobs=config_knobs,
            )
            command.extend(
                [
                    "--random_forrest_execution_time_predictor_config_num_training_job_threads",
                    str(args.predictor_training_job_threads),
                ]
            )
            command = configure_stage_batch_ledger(
                command, enabled=args.store_stage_batch_ledger
            )
            command.extend(
                [
                    "--metrics_config_cache_dir",
                    str(predictor_cache_root),
                ]
            )
            if args.align_real_graph_runtime:
                command.extend(
                    [
                        "--cudagraph_capture_sizes",
                        *(str(size) for size in GRAPH_CAPTURE_SIZES_BY_MNS[config.mns]),
                    ]
                )
            command = BASE.configure_cc_command(
                command,
                backend=args.cc_backend,
                allreduce_csv=args.allreduce_csv,
                cache=args.output_root / "cc-cache",
            )
            BASE.write_json(run_dir / "command.json", command)
            started = time.time()
            with (run_dir / "stdout.log").open("w") as stdout, (
                run_dir / "stderr.log"
            ).open("w") as stderr:
                completed = subprocess.run(
                    command,
                    cwd=args.frontier_source,
                    env=environment,
                    stdout=stdout,
                    stderr=stderr,
                    timeout=args.timeout_seconds,
                    check=False,
                )
            fallback_evidence = (
                collective_fallback_evidence(run_dir)
                if args.cc_backend == "vidur" and config.tp > 1
                else []
            )
            if completed.returncode == 0 and fallback_evidence:
                result = {
                    "status": "frontier_failed",
                    "failure_kind": "collective_profile_fallback",
                    "returncode": completed.returncode,
                    "config": {
                        "tp": config.tp,
                        "mns": config.mns,
                        "name": config.name,
                    },
                    "trace_label": trace["label"],
                    "offered_request_rate": trace["offered_request_rate"],
                    "offered_request_rate_per_gpu": (
                        trace["offered_request_rate"] / config.tp
                    ),
                    "request_count": trace["requests"],
                    "elapsed_seconds": time.time() - started,
                    "trace_sha256": trace["sha256"],
                    "collective_fallback_evidence": fallback_evidence,
                    "allreduce_csv_sha256": allreduce_csv_sha256,
                }
                BASE.write_json(result_path, result)
                loads.append(result)
                if not args.continue_on_failure:
                    raise RuntimeError(
                        f"measured collective fallback for {config.name}/{trace['label']}"
                    )
                continue
            if completed.returncode != 0:
                stderr_path = run_dir / "stderr.log"
                result = {
                    "status": "frontier_failed",
                    "failure_kind": classify_frontier_failure(
                        stderr_path.read_text(errors="replace")
                    ),
                    "returncode": completed.returncode,
                    "config": {
                        "tp": config.tp,
                        "mns": config.mns,
                        "name": config.name,
                    },
                    "trace_label": trace["label"],
                    "offered_request_rate": trace["offered_request_rate"],
                    "offered_request_rate_per_gpu": (
                        trace["offered_request_rate"] / config.tp
                    ),
                    "request_count": trace["requests"],
                    "elapsed_seconds": time.time() - started,
                    "trace_sha256": trace["sha256"],
                    "stderr_sha256": BASE.sha256(stderr_path),
                    "allreduce_csv_sha256": allreduce_csv_sha256,
                }
                BASE.write_json(result_path, result)
                loads.append(result)
                print(
                    json.dumps(
                        {
                            "config": config.name,
                            "trace": trace["label"],
                            "status": result["status"],
                            "failure_kind": result["failure_kind"],
                        },
                        sort_keys=True,
                    ),
                    flush=True,
                )
                if not args.continue_on_failure:
                    raise RuntimeError(
                        f"Frontier failed for {config.name}/{trace['label']}: "
                        f"{completed.returncode}"
                    )
                continue
            metrics = find_request_metrics(run_dir)
            result = {
                "status": "completed",
                "config": {"tp": config.tp, "mns": config.mns, "name": config.name},
                "trace_label": trace["label"],
                "offered_request_rate": trace["offered_request_rate"],
                "offered_request_rate_per_gpu": trace["offered_request_rate"] / config.tp,
                "request_count": trace["requests"],
                "elapsed_seconds": time.time() - started,
                "trace_sha256": trace["sha256"],
                "request_metrics_sha256": BASE.sha256(metrics),
                "score": score(metrics, trace["shapes"]),
                "allreduce_csv_sha256": allreduce_csv_sha256,
            }
            BASE.write_json(result_path, result)
            loads.append(result)
            print(
                json.dumps(
                    {
                        "config": config.name,
                        "trace": trace["label"],
                        "rate": trace["offered_request_rate"],
                        "primary": result["score"]["slos"]["tpot_150ms"],
                    },
                    sort_keys=True,
                ),
                flush=True,
            )
        config_results.append(
            {
                "config": {"tp": config.tp, "mns": config.mns, "name": config.name},
                "loads": loads,
            }
        )

    rankings = {}
    for slo in (f"tpot_{int(value)}ms" for value in TPOT_SLOS_MS):
        records = []
        for item in config_results:
            completed_loads = [
                load for load in item["loads"] if load["status"] == "completed"
            ]
            invalid_loads = [
                load for load in item["loads"] if load["status"] != "completed"
            ]
            feasible = [
                load["offered_request_rate"]
                for load in completed_loads
                if load["score"]["slos"][slo]["feasible"]
            ]
            capacity = max(feasible, default=None)
            records.append(
                {
                    "config": item["config"],
                    "ranking_valid": not invalid_loads,
                    "invalid_loads": [
                        {
                            "trace_label": load["trace_label"],
                            "failure_kind": load.get("failure_kind", "unknown"),
                        }
                        for load in invalid_loads
                    ],
                    "maximum_tested_feasible_request_rate": capacity,
                    "maximum_tested_feasible_request_rate_per_gpu": (
                        capacity / item["config"]["tp"] if capacity is not None else None
                    ),
                    "lower_censored": capacity is None and not invalid_loads,
                    "upper_censored": (
                        capacity == traces[-1]["offered_request_rate"]
                        and not invalid_loads
                    ),
                }
            )
        records.sort(
            key=lambda row: (
                -(
                    row["maximum_tested_feasible_request_rate_per_gpu"]
                    if row["maximum_tested_feasible_request_rate_per_gpu"] is not None
                    else -1
                ),
                row["config"]["name"],
            )
        )
        rankings[slo] = records

    has_invalid_cells = any(
        load["status"] != "completed"
        for item in config_results
        for load in item["loads"]
    )
    if selected != list(BASE.GRID):
        manifest_status = "partial_not_decision_bearing"
    elif has_invalid_cells:
        manifest_status = "frozen_with_invalid_cells"
    else:
        manifest_status = "frozen_before_real"
    manifest = {
        "schema": "frontier-qwen30-exact-trace-surface-v1",
        "status": manifest_status,
        "contract": {
            "window_seconds": (
                WINDOW_SECONDS if args.rate_contract == "trace-window" else None
            ),
            "rate_contract": args.rate_contract,
            "prefix_caching": args.prefix_caching,
            "arrival": "original_trace_timestamp_and_order",
            "input_output": "exact_source_values",
            "store_stage_batch_ledger": args.store_stage_batch_ledger,
            "predictor_cache_root": str(predictor_cache_root),
            "ttft_slo": "1000ms + 1000ms * input_tokens / 8000",
            "tpot_slos_ms": TPOT_SLOS_MS,
            "primary_tpot_slo_ms": 150.0,
            "target_pass_rate": TARGET_PASS_RATE,
            "predictor_training_job_threads": args.predictor_training_job_threads,
            "decode_cuda_graph_mode": args.decode_cuda_graph_mode,
            "align_real_graph_runtime": args.align_real_graph_runtime,
            "fresh_predictor_cache": args.fresh_predictor_cache,
        },
        "frontier": {
            "source": str(args.frontier_source),
            "git_head": frontier_head,
            "git_status_short": subprocess.run(
                ["git", "-C", str(args.frontier_source), "status", "--short"],
                check=True,
                text=True,
                stdout=subprocess.PIPE,
            ).stdout,
        },
        "profiles": {
            "root": str(args.profile_root),
            "coverage": coverage,
            "sha256": {name: BASE.sha256(path) for name, path in paths.items()},
        },
        "kernel_only_profiles": (
            None
            if kernel_paths is None
            else {
                "root": str(args.kernel_profile_root),
                "coverage": kernel_coverage,
                "sha256": {
                    name: BASE.sha256(path) for name, path in kernel_paths.items()
                },
            }
        ),
        "runtime_alignment": {
            "capture_sizes_by_mns": (
                GRAPH_CAPTURE_SIZES_BY_MNS if args.align_real_graph_runtime else None
            ),
            "num_blocks_by_config": (
                {
                    f"tp{tp}_mns{mns}": blocks
                    for (tp, mns), blocks in REAL_NUM_BLOCKS_BY_CONFIG.items()
                }
                if args.align_real_graph_runtime
                else None
            ),
        },
        "collective": {
            "backend": args.cc_backend,
            "allreduce_csv": str(args.allreduce_csv) if args.allreduce_csv else None,
            "allreduce_csv_sha256": allreduce_csv_sha256,
            "measured_profile_validation": allreduce_profile,
        },
        "traces": [trace_manifest_entry(trace) for trace in traces],
        "config_results": config_results,
        "rankings": rankings,
    }
    BASE.write_json(args.output_root / "frontier_trace_surface_frozen.json", manifest)
    print(args.output_root / "frontier_trace_surface_frozen.json")


if __name__ == "__main__":
    main()
