bench: include prompts in results bundle
cda326f verified 7 days ago
55.8 kB
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""FastH3 Omni (Ref2VA) MLX benchmark matrix for Apple Silicon.
Runs the same workload matrix as the DGX Spark benchmark: compositions x
decoders (36L / 8L) x canvases (480p / 768p), each cell being one excluded
seed-42 warmup followed by measured seeds 0, 1, 2 at concurrency 1. Each cell
runs in a fresh subprocess so memory never accumulates between cells. The run
is resumable (re-run the same command), survives a crashing or OOM cell, and
ends by writing latency.csv, benchmark.json, README.md and SEND_BACK.tar.gz.
Quick start (see README.md):
    python3 run_matrix.py --smoke          # ~4 minute install check
    python3 run_matrix.py                  # full matrix, resumable
Pure stdlib plus ffmpeg/ffprobe on PATH. The worker subprocess imports mlx and
fastvideo (run this script with the Python environment that has them).
"""
from __future__ import annotations

import argparse
import datetime as dt
import hashlib
import json
import os
import re
import shutil
import signal
import socket
import subprocess
import sys
import tarfile
import threading
import time
import traceback
from pathlib import Path

SCRIPT_DIR = Path(__file__).resolve().parent
sys.path.insert(0, str(SCRIPT_DIR))
import aggregate as agg  # noqa: E402

GIB = 2**30
WARMUP_SEED = 42
EXPECTED_FRAMES = {5: 124, 10: 243, 15: 362}
DEFAULT_CACHE_LIMIT_GIB = 8.0
WIRED_FRACTION = 0.66
WIRED_CAP_GIB = 64.0
SAMPLE_INTERVAL_S = 2.0
HEAD_HASH_BYTES = 64 * 1024 * 1024
LOCK_DIR = Path("/tmp/fh3mlx/gpu.lock")
LOCK_POLL_S = 30
BUNDLE_MP4_CAP_BYTES = 2 * GIB
DEFAULT_CELL_TIMEOUT_S = 4 * 3600
SMOKE_WIDTH, SMOKE_HEIGHT = 448, 256
DECODER_FILES = {
    "36l": "vae/fasth3-quality-36L-decoder-nvfp4.safetensors",
    "8l": "vae/fasth3-fast-8L-decoder-nvfp4.safetensors",
}
ENCODER_FILE = "vae/lightvae-lite-h3-encoder.safetensors"
TEXT_ENCODER_DIR = "text_encoder-nvfp4"


def utc_now() -> str:
    return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())


def sha256_file(path: Path, limit: int | None = None) -> str:
    digest = hashlib.sha256()
    remaining = limit
    with open(path, "rb") as handle:
        while remaining is None or remaining > 0:
            block = handle.read(1 << 24 if remaining is None else min(1 << 24, remaining))
            if not block:
                break
            digest.update(block)
            if remaining is not None:
                remaining -= len(block)
    return digest.hexdigest()


def write_json(path: Path, payload: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(payload, indent=2, sort_keys=True, default=str) + "\n", encoding="utf-8")
    temporary.replace(path)


def append_jsonl(path: Path, record: dict) -> None:
    """Durably append one record so a crash cannot erase completed requests."""
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "a", encoding="utf-8") as handle:
        handle.write(json.dumps(record, sort_keys=True, default=str) + "\n")
        handle.flush()
        os.fsync(handle.fileno())


def run_text(command: list[str]) -> str:
    try:
        done = subprocess.run(command, capture_output=True, text=True, check=False, timeout=30)
    except (OSError, subprocess.TimeoutExpired):
        return ""
    return done.stdout.strip()


# --------------------------------------------------------------------------- data bundle


def verify_data_bundle(data_root: Path) -> dict:
    """Check every SHA256SUMS entry; returns identity facts and a path->sha map. Raises on any mismatch."""
    sums = data_root / "SHA256SUMS"
    if not sums.is_file():
        raise FileNotFoundError(f"{sums} not found. Point --data-root (or OMNIREF_DATA) at the omniref-handoff bundle.")
    entries: dict[str, str] = {}
    bad: list[str] = []
    for line in sums.read_text(encoding="utf-8").splitlines():
        if not line.strip():
            continue
        digest, _, name = line.partition("  ")
        name = name.strip().lstrip("*")
        path = data_root / name
        entries[name] = digest
        if not path.is_file():
            bad.append(f"missing: {name}")
        elif sha256_file(path) != digest:
            bad.append(f"mismatch: {name}")
    if bad:
        raise ValueError(f"data bundle verification failed ({len(bad)} problems): " + "; ".join(bad[:10]))
    return {
        "root": str(data_root),
        "sha256sums_sha256": sha256_file(sums),
        "sha256sums_lines": len(entries),
        "verification": "SHA256SUMS verified by run_matrix.py",
        "file_sha256": entries,
    }


def load_index(data_root: Path) -> dict:
    return json.loads((data_root / "prompts" / "index.json").read_text(encoding="utf-8"))


def index_record(index: dict, composition: str, seconds: int) -> dict:
    matches = [r for r in index["records"] if r["composition"] == composition and int(r["target_seconds"]) == seconds]
    if len(matches) != 1:
        raise ValueError(f"expected one index record for {composition} {seconds}s, got {len(matches)}")
    record = matches[0]
    if int(record["aligned_frames"]) != EXPECTED_FRAMES[seconds]:
        raise ValueError(f"{composition} {seconds}s has unexpected frame count {record['aligned_frames']}")
    return record


def read_manifest(data_root: Path, relative: str) -> list[dict]:
    path = data_root / relative
    rows = [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
    if not rows:
        raise ValueError(f"empty manifest {path}")
    return rows


def request_rows(data_root: Path, index: dict, composition: str, seconds: int) -> tuple[dict, dict[int, dict]]:
    """(warmup row for seed 42, {seed: measured row}) from the bundle's manifests."""
    record = index_record(index, composition, seconds)
    warmup = read_manifest(data_root, record["warmup_manifest"])
    measured = read_manifest(data_root, record["measured_manifest"])
    if [int(r["seed"]) for r in warmup] != [WARMUP_SEED]:
        raise ValueError(f"{composition} {seconds}s warmup manifest is not seed {WARMUP_SEED}")
    if [int(r["seed"]) for r in measured] != list(agg.DEFAULT_SEEDS):
        raise ValueError(f"{composition} {seconds}s measured seeds are not 0/1/2")
    for row in (*warmup, *measured):
        if int(row["metadata"]["aligned_frames"]) != EXPECTED_FRAMES[seconds]:
            raise ValueError(f"manifest row {row['id']} has inconsistent frame count")
        if row["conditions"] != record["conditions"]:
            raise ValueError(f"manifest row {row['id']} conditions differ from the prompt index")
    return warmup[0], {int(r["seed"]): r for r in measured}


def media_path(data_root: Path, condition: dict) -> Path:
    root = (data_root / "inputs").resolve()
    path = (root / condition["media"]).resolve()
    if root not in path.parents:
        raise ValueError(f"condition media escapes the input root: {path}")
    if not path.is_file():
        raise FileNotFoundError(f"condition media is missing: {path}")
    return path


def condition_entries(row: dict, file_sha256: dict[str, str]) -> list[dict]:
    return [{
        "type": c["type"],
        "media": c["media"],
        "sha256": file_sha256.get(f"inputs/{c['media']}"),
    } for c in row["conditions"]]


# --------------------------------------------------------------------------- model layout


def default_model_dir() -> Path:
    """HF layout: this script lives in <repo>/bench/, the model root in <repo>/. Falls back to the local layout."""
    candidates = [SCRIPT_DIR.parent, Path.home() / "models" / "fasth3-omni-mlx"]
    return next((c for c in candidates if (c / "root").is_dir()), candidates[0])


def resolve_model_paths(model_dir: Path, dit: str, decoder: str, conditioner_override: Path | None) -> dict:
    dit_path = Path(dit).expanduser() if os.sep in dit or Path(dit).expanduser().is_dir() else model_dir / "dit" / dit
    root = model_dir / "root"
    conditioner = conditioner_override
    if conditioner is None and not (root / "text_encoder" / "model.safetensors").exists():
        conditioner = model_dir / TEXT_ENCODER_DIR  # symlink missing or broken after a download/copy
    return {
        "model_root": root,
        "dit": dit_path,
        "vae_decoder": model_dir / DECODER_FILES[decoder],
        "vae_encoder": model_dir / ENCODER_FILE,
        "conditioner_dir": conditioner,
    }


def preflight_paths(paths: dict, decoders: list[str], model_dir: Path, dit: str) -> list[str]:
    problems = []
    root = paths["model_root"]
    for rel in ("fastvideo_inference.json", "tokenizer", "processor", "audio_vae", "scheduler", "vae/config.json"):
        if not (root / rel).exists():
            problems.append(f"{root / rel}")
    if not (paths["dit"] / "mlx_h3_dit.json").is_file():
        problems.append(str(paths["dit"] / "mlx_h3_dit.json"))
    conditioner = paths["conditioner_dir"] or root / "text_encoder"
    if not (Path(conditioner) / "model.safetensors").exists():
        problems.append(str(Path(conditioner) / "model.safetensors"))
    for decoder in decoders:
        if not (model_dir / DECODER_FILES[decoder]).is_file():
            problems.append(str(model_dir / DECODER_FILES[decoder]))
    if not (model_dir / ENCODER_FILE).is_file():
        problems.append(str(model_dir / ENCODER_FILE))
    return problems


def file_identity(path: Path) -> dict:
    path = path.resolve()
    return {"path": str(path), "bytes": path.stat().st_size, "head64mb_sha256": sha256_file(path, HEAD_HASH_BYTES)}


def weights_identity(model_dir: Path, dit_dir: Path, decoders: list[str], conditioner: Path | None) -> dict:
    root = model_dir / "root"
    text_dir = conditioner or root / "text_encoder"
    identity: dict = {
        "dit_dir": str(dit_dir),
        "dit_manifest_sha256": sha256_file(dit_dir / "mlx_h3_dit.json"),
        "dit_weights": [file_identity(p) for p in sorted(dit_dir.glob("*.safetensors"))],
        "contract_sha256": sha256_file(root / "fastvideo_inference.json"),
        "text_encoder": [file_identity(p) for p in sorted(Path(text_dir).resolve().glob("*.safetensors"))],
        "text_encoder_manifest_sha256": None,
        "vae_encoder": file_identity(model_dir / ENCODER_FILE),
        "vae_decoders": {d: file_identity(model_dir / DECODER_FILES[d]) for d in decoders},
    }
    manifest = Path(text_dir) / "mlx_h3_nvfp4_encoder.json"
    if manifest.is_file():
        identity["text_encoder_manifest_sha256"] = sha256_file(manifest)
    return identity


def num_steps_from_contract(model_root: Path) -> int:
    contract = json.loads((model_root / "fastvideo_inference.json").read_text(encoding="utf-8"))
    return int(contract.get("transformer_forwards", contract.get("num_inference_steps", 8)))


# --------------------------------------------------------------------------- ffprobe / output validation


def ffprobe_media(ffprobe: str, path: Path) -> dict:
    command = [ffprobe, "-v", "error", "-show_entries",
               "stream=index,codec_name,codec_type,width,height,r_frame_rate,sample_rate,channels,duration,nb_frames"
               ":format=duration,size,format_name", "-of", "json", str(path)]
    done = subprocess.run(command, capture_output=True, text=True, check=False)
    if done.returncode != 0:
        return {"ok": False, "error": done.stderr.strip()[:2000]}
    try:
        return {"ok": True, "probe": json.loads(done.stdout)}
    except json.JSONDecodeError as error:
        return {"ok": False, "error": f"ffprobe JSON decode failed: {error}"}


def output_probe(ffprobe: str, path: Path, out_dir: Path, expected_frames: int, width: int, height: int) -> dict:
    if not path.is_file():
        return {"ok": False, "error": "output MP4 does not exist", "warnings": []}
    probe = ffprobe_media(ffprobe, path)
    payload: dict = {"ok": False, "bytes": path.stat().st_size, "sha256": sha256_file(path),
                     "relative_path": str(path.relative_to(out_dir)), "warnings": []}
    if not probe["ok"]:
        payload["error"] = probe.get("error", "ffprobe failed")
        return payload
    streams = probe["probe"].get("streams", [])
    video = [s for s in streams if s.get("codec_type") == "video"]
    audio = [s for s in streams if s.get("codec_type") == "audio"]
    payload.update(video_streams=video, audio_streams=audio, format=probe["probe"].get("format", {}))
    if not video:
        payload["error"] = "no video stream"
    elif not audio:
        payload["error"] = "no audio stream"
    else:
        payload["ok"] = True
        frames = video[0].get("nb_frames")
        if frames is not None and int(frames) != expected_frames:
            payload["warnings"].append(f"video has {frames} frames, expected {expected_frames}")
        if (video[0].get("width"), video[0].get("height")) != (width, height):
            payload["warnings"].append(
                f"video is {video[0].get('width')}x{video[0].get('height')}, requested {width}x{height}")
    return payload


# --------------------------------------------------------------------------- GPU lock


class GpuLock:
    """Cross-job GPU lock: atomic mkdir of /tmp/fh3mlx/gpu.lock with our PID in owner. Released in finally."""

    def __init__(self, path: Path = LOCK_DIR, steal_stale: bool = False) -> None:
        self.path = path
        self.steal_stale = steal_stale
        self.held = False

    def _owner(self) -> str:
        try:
            return (self.path / "owner").read_text(encoding="utf-8").strip()
        except OSError:
            return "unknown"

    def _owner_dead(self) -> bool:
        owner = self._owner()
        if not owner.isdigit():
            return False
        try:
            os.kill(int(owner), 0)
        except ProcessLookupError:
            return True
        except PermissionError:
            return False
        return False

    def acquire(self) -> None:
        self.path.parent.mkdir(parents=True, exist_ok=True)
        waited = 0
        while True:
            try:
                os.mkdir(self.path)
                break
            except FileExistsError:
                if self.steal_stale and self._owner_dead():
                    print(f"[lock] owner pid {self._owner()} is dead; removing stale lock", flush=True)
                    shutil.rmtree(self.path, ignore_errors=True)
                    continue
                if waited % (10 * LOCK_POLL_S) == 0:
                    print(f"[lock] GPU lock {self.path} held by pid {self._owner()}; waiting "
                          f"(retry every {LOCK_POLL_S} s, waited {waited} s)", flush=True)
                time.sleep(LOCK_POLL_S)
                waited += LOCK_POLL_S
        self.held = True
        (self.path / "owner").write_text(f"{os.getpid()}\n", encoding="utf-8")
        print(f"[lock] acquired {self.path} (pid {os.getpid()})", flush=True)

    def release(self) -> None:
        if self.held:
            shutil.rmtree(self.path, ignore_errors=True)
            self.held = False
            print(f"[lock] released {self.path}", flush=True)


# --------------------------------------------------------------------------- hardware / receipt


def hardware_facts() -> dict:
    memsize = run_text(["sysctl", "-n", "hw.memsize"])
    return {
        "chip": run_text(["sysctl", "-n", "machdep.cpu.brand_string"]),
        "memory_bytes": int(memsize) if memsize.isdigit() else None,
        "memory_gib": round(int(memsize) / GIB, 2) if memsize.isdigit() else None,
        "macos": run_text(["sw_vers", "-productVersion"]),
        "macos_build": run_text(["sw_vers", "-buildVersion"]),
        "machine": run_text(["uname", "-m"]),
        "cpu_cores": run_text(["sysctl", "-n", "hw.ncpu"]),
        "power": (run_text(["pmset", "-g", "batt"]).splitlines() or [""])[0],
    }


def auto_wired_limit_gib(memory_gib: float | None) -> float | None:
    if not memory_gib:
        return None
    return round(min(WIRED_FRACTION * memory_gib, WIRED_CAP_GIB), 2)


def software_probe_main() -> int:
    """Worker-side: report what the pipeline environment actually is. Exits non-zero if imports fail."""
    import importlib.metadata as metadata
    import mlx.core as mx
    import fastvideo
    from fastvideo.mlx_runtime.minimax_h3_pipeline import MiniMaxH3MLXPipeline  # noqa: F401
    from fastvideo.mlx_runtime.minimax_h3_reference import MiniMaxH3Reference  # noqa: F401

    package_dir = Path(fastvideo.__file__).resolve().parent.parent
    sha = run_text(["git", "-C", str(package_dir), "rev-parse", "HEAD"])
    info_fn = getattr(mx, "device_info", None) or mx.metal.device_info
    print(json.dumps({
        "python": sys.version.split()[0],
        "python_executable": sys.executable,
        "mlx": metadata.version("mlx"),
        "mx_device_info": info_fn(),
        "fastvideo": {
            "file": str(Path(fastvideo.__file__).resolve()),
            "git_sha": sha or None,
            "git_branch": run_text(["git", "-C", str(package_dir), "rev-parse", "--abbrev-ref", "HEAD"]) or None,
            "dirty": bool(run_text(["git", "-C", str(package_dir), "status", "--porcelain", "--untracked-files=no"])),
        },
    }, default=str))
    return 0


def worker_env(args: argparse.Namespace) -> dict:
    env = dict(os.environ, PYTHONUNBUFFERED="1")
    if args.fastvideo_root:
        env["PYTHONPATH"] = os.pathsep.join(filter(None, [str(args.fastvideo_root), env.get("PYTHONPATH", "")]))
    return env


def probe_software(args: argparse.Namespace) -> dict:
    done = subprocess.run([args.python, str(Path(__file__).resolve()), "_software"], capture_output=True, text=True,
                          env=worker_env(args), check=False, cwd=str(SCRIPT_DIR))
    if done.returncode != 0:
        raise RuntimeError("Cannot import mlx / fastvideo with "
                           f"{args.python}. Install them or pass --python / --fastvideo-root.\n{done.stderr[-2000:]}")
    return json.loads(done.stdout.strip().splitlines()[-1])


# --------------------------------------------------------------------------- CLI


def parse_csv(value: str) -> list[str]:
    return [item.strip() for item in value.split(",") if item.strip()]


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--out", type=Path, help="results directory (default ./results-<hostname>-<date>[-smoke])")
    parser.add_argument("--data-root", type=Path,
                        default=Path(os.environ.get("OMNIREF_DATA", Path.home() / "omniref-handoff")),
                        help="omniref-handoff data bundle (env OMNIREF_DATA)")
    parser.add_argument("--model-dir", type=Path, default=Path(os.environ.get("OMNIREF_MODEL", default_model_dir())),
                        help="dir holding root/, dit/<fmt>/, text_encoder-nvfp4/, vae/ (env OMNIREF_MODEL)")
    parser.add_argument("--dit", default="int6", help="DiT conversion under <model-dir>/dit/ (int6, int8, bf16) or a path")
    parser.add_argument("--quant-gemm", choices=("auto", "dequant", "qmm"), default="auto",
                        help="INT8/INT6 linear path (auto|dequant|qmm); ignored by a bf16 DiT")
    parser.add_argument("--m5-sweep", action="store_true",
                        help="fastest-config sweep: image composition, 480p, 36L decoder, seed 0 only, over "
                        "dit {int6,int8,bf16 if present} x quant-gemm {dequant,qmm} (bf16 ignores quant-gemm). "
                        "Writes m5-sweep.md/csv")
    parser.add_argument("--conditioner-dir", type=Path, help="text encoder dir; default root/text_encoder, falling "
                        "back to <model-dir>/text_encoder-nvfp4 when that symlink is missing")
    parser.add_argument("--compositions", type=parse_csv,
                        help="comma list; default image,image-voice,images3,clip,video. Also: mixed,clips3,images9")
    parser.add_argument("--seconds", type=int, default=5, choices=sorted(EXPECTED_FRAMES), help="clip length (default 5)")
    parser.add_argument("--decoders", type=parse_csv, help="36l,8l (default both)")
    parser.add_argument("--canvases", type=parse_csv, help="480p,768p (default both); or WxH")
    parser.add_argument("--width", type=int, help="custom canvas width (needs --height)")
    parser.add_argument("--height", type=int, help="custom canvas height (needs --width)")
    parser.add_argument("--seeds", type=lambda v: [int(x) for x in parse_csv(v)], default=list(agg.DEFAULT_SEEDS),
                        help="measured seeds, subset of 0,1,2 (default 0,1,2)")
    parser.add_argument("--no-warmup", action="store_true", help="skip the excluded seed-42 warmup")
    parser.add_argument("--smoke", action="store_true",
                        help=f"install check: one seed-42 request, 36L, {SMOKE_WIDTH}x{SMOKE_HEIGHT}, composition image")
    parser.add_argument("--vae-dtype", default="bf16")
    parser.add_argument("--metal-wired-limit-gib", default="auto",
                        help=f"GiB to wire, or 'auto' (~{WIRED_FRACTION} x memory, cap {WIRED_CAP_GIB:g}); 'off' to skip")
    parser.add_argument("--mlx-cache-limit-gib", type=float, default=DEFAULT_CACHE_LIMIT_GIB,
                        help="mx.set_cache_limit in GiB (negative = leave unset)")
    parser.add_argument("--reference-max-pixels", default=None,
                        help="cap each reference canvas: a pixel count or 'match'. Default unset = trained canvas")
    parser.add_argument("--steps", type=int, default=0, help="transformer forwards; 0 = contract's count")
    parser.add_argument("--retry-failed", action="store_true", help="re-run cells previously recorded as failed")
    parser.add_argument("--skip-verify", action="store_true", help="do not hash the data bundle against SHA256SUMS")
    parser.add_argument("--cell-timeout-s", type=float, default=DEFAULT_CELL_TIMEOUT_S, help="kill a cell after this long")
    parser.add_argument("--python", default=sys.executable, help="python used for cell subprocesses")
    parser.add_argument("--fastvideo-root", type=Path, default=os.environ.get("FASTVIDEO_ROOT"),
                        help="FastVideo checkout to put on PYTHONPATH (env FASTVIDEO_ROOT)")
    parser.add_argument("--ffprobe", default=os.environ.get("OMNIREF_FFPROBE") or shutil.which("ffprobe") or "ffprobe")
    parser.add_argument("--lock-dir", type=Path, default=LOCK_DIR)
    parser.add_argument("--steal-stale-lock", action="store_true", help="remove the GPU lock if its owner pid is dead")
    parser.add_argument("--no-bundle", action="store_true", help="do not create SEND_BACK.tar.gz")
    parser.add_argument("--dry-run", action="store_true", help="print the plan and exit; no GPU, no lock")
    parser.add_argument("--self-test", action="store_true", help="validate manifests and CSV schema; no GPU")
    parser.add_argument("--aggregate-only", action="store_true", help="rebuild csv/json/README from --out and exit")
    return parser


def resolve_selection(args: argparse.Namespace) -> dict:
    """Turn flags (including --smoke defaults) into the ordered list of cells."""
    if args.m5_sweep:
        return resolve_m5_sweep(args)
    custom = args.width is not None or args.height is not None
    if custom and (args.width is None or args.height is None):
        raise SystemExit("--width and --height must be given together")
    if args.smoke and not custom and not args.canvases:
        args.width, args.height, custom = SMOKE_WIDTH, SMOKE_HEIGHT, True
    compositions = args.compositions or (["image"] if args.smoke else list(agg.DEFAULT_COMPOSITIONS))
    decoders = [d.lower() for d in (args.decoders or (["36l"] if args.smoke else list(agg.DECODERS)))]
    canvases = {}
    labels = [f"{args.width}x{args.height}"] if custom else (args.canvases or list(agg.STANDARD_CANVASES))
    for label in labels:
        if label in agg.STANDARD_CANVASES:
            canvases[label] = agg.STANDARD_CANVASES[label]
        elif re.fullmatch(r"\d+x\d+", label):
            width, height = (int(v) for v in label.split("x"))
            canvases[label] = (width, height)
        else:
            raise SystemExit(f"unknown canvas {label!r}; use 480p, 768p or WxH")
    bad = [d for d in decoders if d not in agg.DECODERS]
    if bad:
        raise SystemExit(f"unknown decoder(s) {bad}; use 36l and/or 8l")
    seeds = [agg.WARMUP_SEED] if args.smoke else args.seeds
    if not args.smoke and not set(seeds) <= set(agg.DEFAULT_SEEDS):
        raise SystemExit("--seeds must be a subset of 0,1,2 (the bundle's measured seeds)")
    cells = []
    for label, (width, height) in canvases.items():  # all 480p first, 36L before 8L, compositions in order
        for decoder in decoders:
            for composition in compositions:
                cells.append({"cell_id": agg.cell_id_for(composition, args.seconds, decoder, label),
                              "composition": composition, "target_seconds": args.seconds, "decoder": decoder,
                              "canvas": label, "width": width, "height": height, "smoke": args.smoke})
    return {"cells": cells, "seeds": seeds, "decoders": decoders, "compositions": compositions}


SWEEP_DITS = ("int6", "int8", "bf16")
SWEEP_GEMMS = ("dequant", "qmm")


def sweep_cell_id(seconds: int, dit: str, gemm: str | None) -> str:
    base = agg.cell_id_for("image", seconds, "36l", "480p")
    return f"{base}-dit-{dit}" + (f"-gemm-{gemm}" if gemm else "")


def resolve_m5_sweep(args: argparse.Namespace) -> dict:
    """Cells for the M5 fastest-config sweep: one cell per (dit, quant-gemm) at image/480p/36L, seed 0."""
    width, height = agg.STANDARD_CANVASES["480p"]
    cells = []
    for dit in SWEEP_DITS:
        if dit == "bf16" and not (args.model_dir / "dit" / "bf16" / "mlx_h3_dit.json").is_file():
            print("[m5-sweep] no dit/bf16 checkpoint found; sweeping int6/int8 only", flush=True)
            continue
        for gemm in ((None,) if dit == "bf16" else SWEEP_GEMMS):
            cells.append({"cell_id": sweep_cell_id(args.seconds, dit, gemm), "composition": "image",
                          "target_seconds": args.seconds,
                          "decoder": "36l", "canvas": "480p", "width": width, "height": height, "smoke": False,
                          "dit": dit, "quant_gemm": gemm or "auto"})
    return {"cells": cells, "seeds": [0], "decoders": ["36l"], "compositions": ["image"]}


def cell_dit(cell: dict, args: argparse.Namespace) -> str:
    return cell.get("dit") or args.dit


def cell_quant_gemm(cell: dict, args: argparse.Namespace) -> str:
    return cell.get("quant_gemm") or args.quant_gemm


def write_sweep_summary(out: Path, cells: list[dict]) -> None:
    """m5-sweep.csv/md: one row per (dit, quant-gemm) cell, ranked by mean seed-0 wall time."""
    records = agg.read_jsonl(out / "per-request-records.jsonl")
    rows = []
    for cell in cells:
        summary = agg.summarize_cell(cell, records, agg.read_json(out / "status" / f"{cell['cell_id']}.json"))
        seed0 = summary["by_seed"].get(0)
        timings = (seed0 or {}).get("timings", {})
        denoise = next((v for k, v in timings.items() if "denoise" in k or "dit" in k), None)
        rows.append({"dit": cell["dit"], "quant_gemm": cell["quant_gemm"] if cell["dit"] != "bf16" else "n/a",
                     "status": summary["classification"], "wall_s": seed0["wall_s"] if seed0 else None,
                     "denoise_s": denoise, "peak_gib": max(((seed0 or {}).get("peak_memory_gib") or {0: 0}).values())})
    rows.sort(key=lambda r: (r["wall_s"] is None, r["wall_s"] or 0.0))
    header = ["dit", "quant_gemm", "status", "wall_s", "denoise_s", "peak_gib"]
    (out / "m5-sweep.csv").write_text("\n".join([",".join(header)] + [",".join(str(r[h]) for h in header) for r in rows]) + "\n")
    lines = ["# M5 fastest-config sweep (image, 480p, 36L, seed 0)", "", "| " + " | ".join(header) + " |",
             "| " + " | ".join("---" for _ in header) + " |"]
    lines += ["| " + " | ".join(str(r[h]) for h in header) + " |" for r in rows]
    (out / "m5-sweep.md").write_text("\n".join(lines) + "\n")


def wired_limit_value(args: argparse.Namespace, memory_gib: float | None) -> float | None:
    value = str(args.metal_wired_limit_gib).lower()
    if value == "auto":
        return auto_wired_limit_gib(memory_gib)
    if value in ("off", "none", "0"):
        return None
    return float(value)


# --------------------------------------------------------------------------- orchestrator


def remaining_work(cell: dict, seeds: list[int], out: Path, args: argparse.Namespace) -> tuple[list[dict], str]:
    """Plan items still to run for a cell, plus a skip reason (empty if the cell must run)."""
    records = agg.read_jsonl(out / "per-request-records.jsonl")
    status = agg.read_json(out / "status" / f"{cell['cell_id']}.json")
    summary = agg.summarize_cell({**cell, "smoke": cell["smoke"]}, records,
                                 {**status, "requested_seeds": seeds, "smoke": cell["smoke"]})
    if summary["classification"] == "measured":
        return [], "already measured"
    if summary["classification"] == "failed" and status.get("state") == "failed" and not args.retry_failed:
        return [], "previously failed (use --retry-failed)"
    todo = [s for s in seeds if s not in summary["by_seed"]]
    plan = []
    if not cell["smoke"] and not args.no_warmup:
        plan.append({"partition": "warmup", "seed": agg.WARMUP_SEED})
    plan += [{"partition": "smoke" if cell["smoke"] else "measured", "seed": s} for s in todo]
    return plan, ""


def tail_lines(path: Path, count: int = 40) -> str:
    try:
        lines = path.read_text(encoding="utf-8", errors="replace").splitlines()
    except OSError:
        return ""
    return "\n".join(lines[-count:])[-4000:]


def terminate_group(proc: subprocess.Popen) -> None:
    if proc.poll() is not None:
        return
    for sig, grace in ((signal.SIGTERM, 20), (signal.SIGKILL, 10)):
        try:
            os.killpg(proc.pid, sig)
        except ProcessLookupError:
            return
        try:
            proc.wait(timeout=grace)
            return
        except subprocess.TimeoutExpired:
            continue


def run_cell_subprocess(spec_path: Path, log_path: Path, args: argparse.Namespace, cell_id: str) -> tuple[int, bool]:
    """Run the worker; returns (returncode, timed_out). Always reaps the child, even on Ctrl-C."""
    log_path.parent.mkdir(parents=True, exist_ok=True)
    started = time.time()
    with open(log_path, "ab") as log:
        proc = subprocess.Popen([args.python, str(Path(__file__).resolve()), "_worker", str(spec_path)], stdout=log,
                                stderr=subprocess.STDOUT, env=worker_env(args), start_new_session=True,
                                cwd=str(SCRIPT_DIR))
        try:
            while True:
                try:
                    return proc.wait(timeout=60), False
                except subprocess.TimeoutExpired:
                    elapsed = time.time() - started
                    last = (tail_lines(log_path, 1) or "").strip()[:140]
                    print(f"[{cell_id}] running {elapsed / 60:.0f} min | {last}", flush=True)
                    if elapsed > args.cell_timeout_s:
                        terminate_group(proc)
                        return proc.returncode if proc.returncode is not None else -9, True
        finally:
            terminate_group(proc)


def record_cell_crash(out: Path, cell: dict, returncode: int, timed_out: bool, log_path: Path, seeds: list[int]) -> None:
    sig = -returncode if returncode < 0 else None
    note = "timeout" if timed_out else (f"killed by signal {sig}" if sig else f"exit code {returncode}")
    hint = " (likely out of memory)" if sig in (6, 9) else ""
    append_jsonl(out / "per-request-records.jsonl", {
        "schema": agg.RECORD_SCHEMA, "cell_id": cell["cell_id"], "composition": cell["composition"],
        "target_seconds": cell["target_seconds"], "decoder": cell["decoder"], "canvas": cell["canvas"],
        "status": "failed", "error_type": "CellProcessDied", "error_message": f"cell process {note}{hint}",
        "exit_code": returncode, "signal": sig, "log_tail": tail_lines(log_path), "finished_utc": utc_now()})
    status = agg.read_json(out / "status" / f"{cell['cell_id']}.json")
    write_json(out / "status" / f"{cell['cell_id']}.json",
               {**status, "state": "failed", "finished_utc": utc_now(), "failure": f"cell process {note}{hint}"})


def make_bundle(out: Path) -> Path:
    bundle = out / "SEND_BACK.tar.gz"
    files = [p for p in out.rglob("*") if p.is_file() and p.name != bundle.name and not p.name.endswith(".tmp")]
    total = sum(p.stat().st_size for p in files)
    drop_mp4 = total > BUNDLE_MP4_CAP_BYTES
    keep = re.compile(r"-(seed0|smoke)\.mp4$")
    temporary = out / "SEND_BACK.tar.gz.tmp"
    with tarfile.open(temporary, "w:gz") as tar:
        for path in sorted(files):
            if drop_mp4 and path.suffix == ".mp4" and not keep.search(path.name):
                continue
            tar.add(path, arcname=f"{out.name}/{path.relative_to(out)}")
    temporary.replace(bundle)
    note = " (MP4s trimmed to one seed-0/smoke video per cell)" if drop_mp4 else ""
    print(f"[bundle] {bundle} {bundle.stat().st_size / 1e6:.1f} MB{note}", flush=True)
    return bundle


def write_receipt(out: Path, args: argparse.Namespace, software: dict, data: dict, weights: dict, settings: dict) -> None:
    path = out / "environment_receipt.json"
    previous = agg.read_json(path)
    hardware = hardware_facts()
    write_json(path, {
        "created_utc": previous.get("created_utc", utc_now()),
        "updated_utc": utc_now(),
        "hostname": socket.gethostname(),
        "hardware": hardware,
        "software": {"python": software["python"], "mlx": software["mlx"], "fastvideo": software["fastvideo"],
                     "mx_device_info": software["mx_device_info"], "python_executable": software["python_executable"],
                     "ffmpeg": (run_text(["ffmpeg", "-version"]).splitlines() or [""])[0],
                     "ffprobe": (run_text([args.ffprobe, "-version"]).splitlines() or [""])[0]},
        "data": {k: v for k, v in data.items() if k != "file_sha256"},
        "weights": weights,
        "run_settings": settings,
        "invocations": [*previous.get("invocations", []), {"utc": utc_now(), "argv": sys.argv}],
        "notes": "Column names nvfp4_* keep the Spark spelling; on Apple Silicon the NVFP4 decoders run as W4A16.",
    })


def orchestrate(args: argparse.Namespace) -> int:
    selection = resolve_selection(args)
    if args.smoke and args.reference_max_pixels is None:
        args.reference_max_pixels = "match"  # references at the trained 1344x768 canvas dominate a tiny smoke run
        print("[smoke] --reference-max-pixels defaults to 'match' (the smoke canvas area)", flush=True)
    cells, seeds = selection["cells"], selection["seeds"]
    hardware = hardware_facts()
    out = (args.out or Path(f"results-{socket.gethostname().split('.')[0]}-{dt.date.today():%Y%m%d}"
                            f"{'-smoke' if args.smoke else ''}")).resolve()
    if args.aggregate_only:
        agg.aggregate(out, cells)
        print(f"aggregated {out}")
        return 0
    wired = wired_limit_value(args, hardware["memory_gib"])
    print(f"[plan] {len(cells)} cells -> {out}", flush=True)
    for cell in cells:
        print(f"  {cell['cell_id']}  {cell['width']}x{cell['height']}", flush=True)
    problems = []
    for dit in dict.fromkeys(cell_dit(cell, args) for cell in cells):
        problems += preflight_paths(resolve_model_paths(args.model_dir, dit, selection["decoders"][0], args.conditioner_dir),
                                    selection["decoders"], args.model_dir, dit)
    problems = list(dict.fromkeys(problems))
    if problems:
        raise SystemExit("Missing model files (check --model-dir / --dit):\n  " + "\n  ".join(problems))
    for tool in ("ffmpeg", args.ffprobe):
        if shutil.which(tool) is None:
            raise SystemExit(f"{tool} not found on PATH (install ffmpeg, e.g. `brew install ffmpeg`).")
    if args.dry_run:
        return 0
    data = {"root": str(args.data_root), "file_sha256": {}, "verification": "skipped (--skip-verify)"}
    if args.skip_verify:
        print("[data] verification skipped", flush=True)
    else:
        print("[data] verifying SHA256SUMS ...", flush=True)
        data = verify_data_bundle(args.data_root)
        print(f"[data] ok: {data['sha256sums_lines']} files", flush=True)
    index = load_index(args.data_root)
    for cell in cells:  # fail on a bad manifest now, not hours in
        request_rows(args.data_root, index, cell["composition"], cell["target_seconds"])
    software = probe_software(args)
    contract_steps = args.steps or num_steps_from_contract(args.model_dir / "root")
    reference_paths = resolve_model_paths(args.model_dir, cell_dit(cells[0], args), selection["decoders"][0],
                                          args.conditioner_dir)
    weights = weights_identity(args.model_dir, reference_paths["dit"], selection["decoders"],
                               reference_paths["conditioner_dir"])
    settings = {"dit": reference_paths["dit"].name, "dit_path": str(reference_paths["dit"]), "vae_dtype": args.vae_dtype,
                "metal_wired_limit_gib": wired, "mlx_cache_limit_gib": args.mlx_cache_limit_gib,
                "reference_max_pixels": args.reference_max_pixels, "num_steps": contract_steps,
                "seeds": seeds, "warmup": not args.no_warmup and not args.smoke, "smoke": args.smoke,
                "conditioner_dir": str(reference_paths["conditioner_dir"]) if reference_paths["conditioner_dir"] else None,
                "quant_gemm": args.quant_gemm, "m5_sweep": args.m5_sweep}
    out.mkdir(parents=True, exist_ok=True)
    write_receipt(out, args, software, data, weights, settings)
    # Ship the exact prompts/conditions with the results so the bundle is self-describing.
    prompts_out = out / "prompts"
    prompts_out.mkdir(exist_ok=True)
    for name in ("index.json", "index.csv"):
        source = Path(args.data_root) / "prompts" / name
        if source.is_file():
            shutil.copy2(source, prompts_out / name)
    exit_code = 0
    lock = GpuLock(args.lock_dir, args.steal_stale_lock)
    # SIGINT is explicit: it is ignored by default when launched as a background job (nohup, agents).
    previous_handlers = {s: signal.signal(s, _raise_interrupt) for s in (signal.SIGINT, signal.SIGTERM, signal.SIGHUP)}
    current: dict | None = None
    try:
        lock.acquire()
        for position, cell in enumerate(cells, 1):
            plan, skip = remaining_work(cell, seeds, out, args)
            if skip:
                print(f"[{cell['cell_id']}] skip: {skip} ({position}/{len(cells)})", flush=True)
                continue
            current = cell
            exit_code |= run_one_cell(cell, plan, seeds, out, args, settings, data, weights, position, len(cells))
            current = None
            agg.aggregate(out, cells)
            if args.m5_sweep:
                write_sweep_summary(out, cells)
    except KeyboardInterrupt:
        print("\n[interrupted] stopping; releasing GPU lock. Re-run the same command to resume.", flush=True)
        if current is not None:
            status = agg.read_json(out / "status" / f"{current['cell_id']}.json")
            write_json(out / "status" / f"{current['cell_id']}.json", {**status, "state": "interrupted"})
        exit_code = 130
    finally:
        lock.release()
        for sig, handler in previous_handlers.items():
            signal.signal(sig, handler)
        agg.aggregate(out, cells)
        if args.m5_sweep:
            write_sweep_summary(out, cells)
        if not args.no_bundle:
            make_bundle(out)
    print(f"[done] results in {out}", flush=True)
    return exit_code


def _raise_interrupt(signum, frame) -> None:  # noqa: ARG001
    raise KeyboardInterrupt


def run_one_cell(cell: dict, plan: list[dict], seeds: list[int], out: Path, args: argparse.Namespace, settings: dict,
                 data: dict, weights: dict, position: int, total: int) -> int:
    cell_id = cell["cell_id"]
    paths = resolve_model_paths(args.model_dir, cell_dit(cell, args), cell["decoder"], args.conditioner_dir)
    settings = {**settings, "dit": paths["dit"].name, "dit_path": str(paths["dit"]),
                "quant_gemm": cell_quant_gemm(cell, args)}
    spec = {
        "cell": cell, "plan": plan, "requested_seeds": seeds, "out": str(out), "data_root": str(args.data_root),
        "ffprobe": args.ffprobe, "settings": settings, "file_sha256": data["file_sha256"],
        "paths": {k: (str(v) if v is not None else None) for k, v in paths.items()},
        "weights": {"dit_manifest_sha256": sha256_file(paths["dit"] / "mlx_h3_dit.json"),
                    "decoder": weights["vae_decoders"][cell["decoder"]], "encoder": weights["vae_encoder"]},
    }
    spec_path = out / "status" / f"{cell_id}.spec.json"
    write_json(spec_path, spec)
    write_json(out / "status" / f"{cell_id}.json", {
        "cell_id": cell_id, "composition": cell["composition"], "target_seconds": cell["target_seconds"],
        "decoder": cell["decoder"], "canvas": cell["canvas"], "width": cell["width"], "height": cell["height"],
        "smoke": cell["smoke"], "requested_seeds": seeds, "state": "starting", "started_utc": utc_now(),
        "dit": settings["dit"], "quant_gemm": settings["quant_gemm"]})
    print(f"[{cell_id}] start ({position}/{total}); plan: "
          + ", ".join(f"{p['partition']}:{p['seed']}" for p in plan), flush=True)
    log_path = out / "logs" / f"{cell_id}.log"
    code, timed_out = run_cell_subprocess(spec_path, log_path, args, cell_id)
    status = agg.read_json(out / "status" / f"{cell_id}.json")
    if status.get("state") not in ("complete", "failed"):
        record_cell_crash(out, cell, code, timed_out, log_path, seeds)
        print(f"[{cell_id}] FAILED: process exited {code}{' (timeout)' if timed_out else ''}", flush=True)
        return 1
    print(f"[{cell_id}] {status['state'].upper()}", flush=True)
    return 0 if status["state"] == "complete" else 1


# --------------------------------------------------------------------------- worker (one cell, fresh process)


class MemorySampler:
    """Every SAMPLE_INTERVAL_S: MLX active/cache/peak GiB and swap used, appended to a JSONL trace."""

    def __init__(self, mx, trace_path: Path) -> None:
        self.mx, self.trace_path = mx, trace_path
        self.maxima = {"max_active_gib": 0.0, "max_cache_gib": 0.0, "max_peak_gib": 0.0, "max_swap_used_mb": None}
        self.first_swap_mb: float | None = None
        self._stop = threading.Event()
        self._thread = threading.Thread(target=self._run, daemon=True)

    @staticmethod
    def swap_used_mb() -> float | None:
        match = re.search(r"used = ([\d.]+)([MGK])", run_text(["sysctl", "-n", "vm.swapusage"]))
        if not match:
            return None
        return float(match.group(1)) * {"K": 1 / 1024, "M": 1, "G": 1024}[match.group(2)]

    def _sample(self, handle) -> None:
        mx = self.mx
        row = {"t": round(time.time(), 2), "active_gib": mx.get_active_memory() / GIB,
               "cache_gib": mx.get_cache_memory() / GIB, "peak_gib": mx.get_peak_memory() / GIB,
               "swap_used_mb": self.swap_used_mb()}
        for key, value in (("max_active_gib", row["active_gib"]), ("max_cache_gib", row["cache_gib"]),
                           ("max_peak_gib", row["peak_gib"])):
            self.maxima[key] = max(self.maxima[key], value)
        if row["swap_used_mb"] is not None:
            self.first_swap_mb = row["swap_used_mb"] if self.first_swap_mb is None else self.first_swap_mb
            self.maxima["max_swap_used_mb"] = max(self.maxima["max_swap_used_mb"] or 0.0, row["swap_used_mb"])
        handle.write(json.dumps({k: (round(v, 3) if isinstance(v, float) else v) for k, v in row.items()}) + "\n")
        handle.flush()

    def _run(self) -> None:
        with open(self.trace_path, "w", encoding="utf-8") as handle:
            while not self._stop.is_set():
                try:
                    self._sample(handle)
                except Exception:  # noqa: BLE001 - a sampling glitch must never kill the benchmark
                    pass
                self._stop.wait(SAMPLE_INTERVAL_S)

    def __enter__(self) -> "MemorySampler":
        self._thread.start()
        return self

    def __exit__(self, *exc_info) -> None:
        self._stop.set()
        self._thread.join(timeout=10)

    def summary(self) -> dict:
        swap_delta = None
        if self.first_swap_mb is not None and self.maxima["max_swap_used_mb"] is not None:
            swap_delta = self.maxima["max_swap_used_mb"] - self.first_swap_mb
        return {**{k: (round(v, 3) if isinstance(v, float) else v) for k, v in self.maxima.items()},
                "swap_growth_mb": swap_delta, "sample_interval_s": SAMPLE_INTERVAL_S}


def reference_max_pixels(value: str | None, width: int, height: int) -> int | None:
    if value is None:
        return None
    if value == "match":
        return width * height
    pixels = int(value)
    if pixels <= 0:
        raise ValueError("--reference-max-pixels must be positive or 'match'")
    return pixels


def _record_base(spec: dict, row: dict, item: dict, tag: str, rewarm: bool, index_rec: dict) -> dict:
    cell, settings = spec["cell"], spec["settings"]
    return {
        "schema": agg.RECORD_SCHEMA, "cell_id": cell["cell_id"], "composition": cell["composition"],
        "target_seconds": cell["target_seconds"], "decoder": cell["decoder"], "canvas": cell["canvas"],
        "partition": item["partition"], "seed": item["seed"], "request_id": row["id"], "tag": tag, "rewarm": rewarm,
        "width": cell["width"], "height": cell["height"], "num_frames": EXPECTED_FRAMES[cell["target_seconds"]],
        "prompt_sha256": index_rec["prompt_sha256"],
        "prompt_sha256_computed": hashlib.sha256(row["prompt"].encode("utf-8")).hexdigest(),
        "condition_media": condition_entries(row, spec["file_sha256"]),
        "settings": settings, "weights": spec["weights"], "model_paths": spec["paths"],
    }


def worker_main(spec_path: str) -> int:
    spec = json.loads(Path(spec_path).read_text(encoding="utf-8"))
    cell, out, settings = spec["cell"], Path(spec["out"]), spec["settings"]
    data_root, cell_id = Path(spec["data_root"]), cell["cell_id"]
    status_path = out / "status" / f"{cell_id}.json"
    records_path = out / "per-request-records.jsonl"
    media_dir = out / "media" / cell_id
    media_dir.mkdir(parents=True, exist_ok=True)
    base_status = agg.read_json(status_path)

    def set_status(**fields) -> None:
        write_json(status_path, {**agg.read_json(status_path), **fields})

    def fail(stage: str, error: BaseException, extra: dict | None = None) -> None:
        append_jsonl(records_path, {
            "schema": agg.RECORD_SCHEMA, "cell_id": cell_id, "composition": cell["composition"],
            "target_seconds": cell["target_seconds"], "decoder": cell["decoder"], "canvas": cell["canvas"],
            "status": "failed", "stage": stage, "error_type": type(error).__name__, "error_message": str(error)[:4000],
            "traceback": traceback.format_exc()[-8000:], "settings": settings, "finished_utc": utc_now(), **(extra or {})})
        print(f"[{cell_id}] FAILED during {stage}: {type(error).__name__}: {error}", flush=True)

    set_status(state="loading_model", pid=os.getpid())
    pipeline = None
    ok = True
    try:
        import mlx.core as mx
        import fastvideo
        from fastvideo.mlx_runtime.minimax_h3_pipeline import MiniMaxH3MLXPipeline
        from fastvideo.mlx_runtime.minimax_h3_reference import MiniMaxH3Reference

        index = load_index(data_root)
        index_rec = index_record(index, cell["composition"], cell["target_seconds"])
        warmup_row, measured_rows = request_rows(data_root, index, cell["composition"], cell["target_seconds"])
        paths = spec["paths"]
        started = time.perf_counter()
        pipeline = MiniMaxH3MLXPipeline(
            model_root=paths["model_root"], mlx_dit_checkpoint=paths["dit"], vae_dtype=settings["vae_dtype"],
            vae_decoder_path=paths["vae_decoder"], vae_encoder_path=paths["vae_encoder"],
            conditioner_dir=paths["conditioner_dir"], metal_wired_limit_gib=settings["metal_wired_limit_gib"],
            mlx_cache_limit_gib=(settings["mlx_cache_limit_gib"] if settings["mlx_cache_limit_gib"] >= 0 else None),
            quant_gemm=settings.get("quant_gemm", "auto"))
        init_s = time.perf_counter() - started
        print(f"[{cell_id}] pipeline ready in {init_s:.1f} s ({fastvideo.__file__})", flush=True)
    except Exception as error:  # noqa: BLE001 - the native failure is the evidence
        fail("setup", error)
        set_status(state="failed", finished_utc=utc_now())
        return 3

    previous = [r for r in agg.read_jsonl(records_path) if r.get("cell_id") == cell_id and r.get("partition") == "warmup"]
    try:
        for item in spec["plan"]:
            ok = run_request(spec, item, pipeline, mx, MiniMaxH3Reference, index_rec, warmup_row, measured_rows, init_s,
                             rewarm=item["partition"] == "warmup" and bool(previous), set_status=set_status) and ok
            if not ok:
                break
    except Exception as error:  # noqa: BLE001
        ok = False
        fail("generate", error)
    finally:
        try:
            pipeline.close()
        except Exception:  # noqa: BLE001 - close must not mask the result
            traceback.print_exc()
    set_status(state="complete" if ok else "failed", finished_utc=utc_now(), model_init_s=round(init_s, 3),
               **{k: v for k, v in base_status.items() if k in ("requested_seeds",)})
    print(f"[{cell_id}] {'COMPLETE' if ok else 'CELL_FAILED'}", flush=True)
    return 0 if ok else 3


def run_request(spec, item, pipeline, mx, reference_cls, index_rec, warmup_row, measured_rows, init_s, rewarm,
                set_status) -> bool:
    cell, settings, out = spec["cell"], spec["settings"], Path(spec["out"])
    data_root = Path(spec["data_root"])
    seed = item["seed"]
    row = warmup_row if seed == agg.WARMUP_SEED else measured_rows[seed]
    tag = {"warmup": "warmup", "smoke": "smoke"}.get(item["partition"], f"seed{seed}")
    output_path = out / "media" / cell["cell_id"] / f"{cell['cell_id']}-{tag}.mp4"
    record = _record_base(spec, row, item, tag, rewarm, index_rec)
    references = [reference_cls(source=str(media_path(data_root, c)), media_type=c["type"]) for c in row["conditions"]]
    set_status(state="request_running", request=row["id"], partition=item["partition"], seed=seed)
    trace_path = out / "media" / cell["cell_id"] / f"{cell['cell_id']}-{tag}.mem.jsonl"
    num_frames = EXPECTED_FRAMES[cell["target_seconds"]]
    try:
        with MemorySampler(mx, trace_path) as sampler:
            started = time.perf_counter()
            result = pipeline.generate(
                row["prompt"], output_path=output_path, references=references, height=cell["height"],
                width=cell["width"], num_frames=num_frames, seed=seed, num_steps=settings["num_steps"],
                reference_max_pixels=reference_max_pixels(settings["reference_max_pixels"], cell["width"], cell["height"]))
            wall_s = time.perf_counter() - started
    except Exception as error:  # noqa: BLE001
        append_jsonl(out / "per-request-records.jsonl", {
            **record, "status": "failed", "stage": "generate", "error_type": type(error).__name__,
            "error_message": str(error)[:4000], "traceback": traceback.format_exc()[-8000:],
            "finished_utc": utc_now()})
        print(f"[{cell['cell_id']}] {row['id']} FAILED: {type(error).__name__}: {error}", flush=True)
        return False
    probe = output_probe(spec["ffprobe"], output_path, out, num_frames, cell["width"], cell["height"])
    record.update(
        wall_s=round(wall_s, 6), timings={k: round(float(v), 4) for k, v in result.timings.items()},
        peak_memory_gib={k: round(float(v), 3) for k, v in result.peak_memory_gib.items()},
        vsa=result.vsa, video_decode_backend=getattr(result, "video_decode_backend", None),
        audio_samples=int(result.waveform.shape[-1]) if getattr(result, "waveform", None) is not None else None,
        pipeline_init_s=round(init_s, 3), memory=sampler.summary(), mem_trace=str(trace_path.relative_to(out)),
        output=probe, validation_warnings=probe.get("warnings", []), status="ok" if probe["ok"] else "output_invalid",
        finished_utc=utc_now())
    append_jsonl(out / "per-request-records.jsonl", record)
    print(f"[{cell['cell_id']}] {row['id']} wall={wall_s:.2f}s peak={max(record['peak_memory_gib'].values(), default=0):.1f}GiB "
          f"swap_max={record['memory']['max_swap_used_mb']}MB ok={probe['ok']} warn={probe.get('warnings')}", flush=True)
    return probe["ok"]


# --------------------------------------------------------------------------- self test


def self_test(args: argparse.Namespace) -> int:
    import tempfile

    problems = []
    sums = args.data_root / "SHA256SUMS"
    if sums.is_file():
        index = load_index(args.data_root)
        count = 0
        for rec in index["records"]:
            request_rows(args.data_root, index, rec["composition"], int(rec["target_seconds"]))
            count += 1
        print(f"manifests ok: {count} index records parsed")
    else:
        print(f"(no data bundle at {args.data_root}; skipping manifest parse)")
    with tempfile.TemporaryDirectory() as tmp:
        out = Path(tmp)
        cell = {"cell_id": "image-5s-36l-480p", "composition": "image", "target_seconds": 5, "decoder": "36l",
                "canvas": "480p", "width": 832, "height": 480, "smoke": False}
        write_json(out / "status" / "image-5s-36l-480p.json", {**cell, "state": "complete", "requested_seeds": [0, 1, 2]})
        for seed, wall in ((0, 10.0), (1, 12.0), (2, 14.0)):
            append_jsonl(out / "per-request-records.jsonl", {
                "cell_id": cell["cell_id"], "partition": "measured", "seed": seed, "status": "ok", "wall_s": wall,
                "output": {"ok": True}, "peak_memory_gib": {"denoise": 20.0}, "memory": {"max_swap_used_mb": 0.0}})
        agg.aggregate(out, [cell])
        header = (out / "latency.csv").read_text(encoding="utf-8").splitlines()[0].split(",")
        expected = ["composition", "target_seconds", "gpu_count"] + [
            f"nvfp4_{d}_{c}_{s}" for d, c in agg.SPARK_COLUMN_ORDER for s in ("mean_s", "sample_std_s", "status")]
        if header != expected:
            problems.append(f"csv header mismatch: {header}")
        row = (out / "latency.csv").read_text(encoding="utf-8").splitlines()[1].split(",")
        if not (row[3] == "12.0" and abs(float(row[4]) - 2.0) < 1e-9 and row[5] == "measured"):
            problems.append(f"csv row mismatch: {row}")
    print("csv schema ok" if not problems else "\n".join(problems))
    return 1 if problems else 0


def main() -> int:
    if len(sys.argv) >= 2 and sys.argv[1] == "_worker":
        return worker_main(sys.argv[2])
    if len(sys.argv) >= 2 and sys.argv[1] == "_software":
        return software_probe_main()
    args = build_parser().parse_args()
    if args.self_test:
        return self_test(args)
    return orchestrate(args)


if __name__ == "__main__":
    sys.exit(main())