scripts/bench_wave_cache_real.py - master - Mihaham/CNN-NEAT

scripts/bench_wave_cache_real.py
815 строк27 KB
MihahamYT
feat(scripts): OVA holdout, hash studies, and cache benches
04 авг 2026, 20:36
04 авг 2026, 20:36e256d0f
#!/usr/bin/env python3
"""Real GPU capacity / throughput test for planned forward wave-cache.

Unlike ``cache_wave_sim`` (fake tensors), this drives real
``ConvolutionalNeuralNetwork`` forwards with CIFAR-10 images and a random
activation (default pool includes GELU), using
``plan_population_misses`` + ``execute_planned_population``.

Measures how many genomes per generation and how many images per eval batch
fit under VRAM, then confirms with several short runs (default 3 x 100 gens).

Usage
-----
::

    # Probe capacity (2 gens each step), then 3 x 100 gens at ~80% of frontier
    python scripts/bench_wave_cache_real.py

    # Fixed size, gelu only, 3 seeds x 100 gens
    python scripts/bench_wave_cache_real.py --no-probe --activation gelu \\
        --pop 64 --images 512 --runs 3 --gens 100

    # Probe only
    python scripts/bench_wave_cache_real.py --probe-only

    # Random activation per run
    python scripts/bench_wave_cache_real.py --activation random --runs 4
"""
from __future__ import annotations

import argparse
import json
import random
import sys
import time
import traceback
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple

import torch
from torchvision import datasets, transforms

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from cnn_neat.binary_head import (  # noqa: E402
    output_to_score,
    predict_from_scores,
    to_binary_labels,
    to_neat_batch,
)
from cnn_neat.config import ALLOWED_ACTIVATIONS  # noqa: E402
from cnn_neat.forward_cache import (  # noqa: E402
    ForwardActivationCache,
    data_id_from_tensor,
)
from cnn_neat.mutations import ConvolutionalNetworkMutator, MutationType  # noqa: E402
from cnn_neat.network import ConvolutionalNeuralNetwork  # noqa: E402

GenomeId = Tuple[int, int]
PopRec = Tuple[ConvolutionalNeuralNetwork, GenomeId]


@dataclass
class GenMetrics:
    generation: int
    n_genomes: int
    n_images: int
    mean_nodes: float
    mean_edges: float
    hit_slots: int
    miss_slots: int
    unique_misses: int
    dedup_saved: int
    hit_rate: float
    wall_s: float
    peak_vram_mb: float
    cache_peak_mb: float
    cache_bytes_mb: float
    elite_fitness: float


@dataclass
class RunResult:
    seed: int
    activation: str
    n_genomes: int
    n_images: int
    generations: int
    ok: bool
    fail_reason: str = ""
    wall_s: float = 0.0
    peak_vram_mb: float = 0.0
    mean_unique_misses: float = 0.0
    mean_hit_rate: float = 0.0
    mean_wall_per_gen_s: float = 0.0
    late_mean_nodes: float = 0.0
    late_mean_edges: float = 0.0
    gens: List[GenMetrics] = field(default_factory=list)


@dataclass
class CapacityPoint:
    n_genomes: int
    n_images: int
    ok: bool
    wall_s: float = 0.0
    peak_vram_mb: float = 0.0
    fail_reason: str = ""


def _sync(device: torch.device) -> None:
    if device.type == "cuda":
        torch.cuda.synchronize(device)


def _peak_vram_mb(device: torch.device) -> float:
    if device.type != "cuda":
        return 0.0
    return float(torch.cuda.max_memory_allocated(device)) / (1024.0 * 1024.0)


def _device_total_mb(device: torch.device) -> float:
    if device.type != "cuda":
        return 0.0
    return float(torch.cuda.get_device_properties(device).total_memory) / (1024.0 * 1024.0)


def _reset_peak(device: torch.device) -> None:
    if device.type == "cuda":
        torch.cuda.reset_peak_memory_stats(device)
        torch.cuda.empty_cache()


class VramBudgetExceeded(RuntimeError):
    """Peak allocated CUDA memory exceeded the configured onboard-VRAM budget."""


def _pick_activation(spec: str, rng: random.Random) -> str:
    if spec == "random":
        return rng.choice(list(ALLOWED_ACTIVATIONS))
    if spec not in ALLOWED_ACTIVATIONS and spec is not None:
        raise ValueError(
            f"Unknown activation {spec!r}. Use one of {ALLOWED_ACTIVATIONS} or 'random'."
        )
    return spec


def _seed_genome(activation: str, rng: random.Random) -> ConvolutionalNeuralNetwork:
    g = ConvolutionalNeuralNetwork(
        input_image_size=(32, 32),
        output_image_size=(8, 8),
        num_input_views=1,
        minimal_connections=True,
        activation=activation,
        edge_weights=False,
        allow_variable_channels=True,
        max_hidden_layers=32,
        max_neurons_per_layer=32,
    )
    # Shared stem so planned-wave overlap is realistic but activations are non-trivial.
    for _ in range(8):
        keys = [k for k, c in g.connections.items() if c.enabled]
        if not keys:
            break
        hid = g.add_node_split(keys[rng.randrange(len(keys))])
        if hid == -1:
            break
    for _ in range(6):
        g = ConvolutionalNetworkMutator(g).apply(MutationType.WITHOUT_SPLIT)
    return g


def _count_edges(g: ConvolutionalNeuralNetwork) -> int:
    return sum(1 for c in g.connections.values() if c.enabled)


def _count_nodes(g: ConvolutionalNeuralNetwork) -> int:
    return int(len(g.all_nodes))


# Bias toward structural change so 100 gens still grow under cache pressure.
_MUTATION_BAG = (
    MutationType.CLONE,
    MutationType.CLONE,
    MutationType.MUTATE_WEIGHTS,
    MutationType.MUTATE_WEIGHTS,
    MutationType.SPLIT,
    MutationType.SPLIT,
    MutationType.WITHOUT_SPLIT,
    MutationType.WITHOUT_SPLIT,
    MutationType.REMOVE_EDGE,
)


def _breed(
    elites: Sequence[ConvolutionalNeuralNetwork],
    *,
    pop_size: int,
    generation: int,
    rng: random.Random,
) -> List[PopRec]:
    assert elites
    out: List[PopRec] = []
    for i, e in enumerate(elites):
        out.append((e, (generation, i)))
    while len(out) < pop_size:
        parent = elites[rng.randrange(len(elites))]
        op = _MUTATION_BAG[rng.randrange(len(_MUTATION_BAG))]
        # Mutator clones internally and returns the (possibly mutated) child.
        child = ConvolutionalNetworkMutator(parent).apply(op)
        out.append((child, (generation, len(out))))
    return out


def load_cifar_neat_batch(
    *,
    n_images: int,
    device: torch.device,
    cifar_root: Path,
    positive_class: int = 0,
    seed: int = 0,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Return ``(neat_images[N,1,3,32,32], binary_labels[N])`` on device."""
    tfm = transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize(
                (0.4914, 0.4822, 0.4465),
                (0.2470, 0.2435, 0.2616),
            ),
        ]
    )
    ds = datasets.CIFAR10(root=str(cifar_root), train=True, download=True, transform=tfm)
    g = torch.Generator().manual_seed(int(seed))
    idx = torch.randperm(len(ds), generator=g)[: int(n_images)].tolist()
    xs: List[torch.Tensor] = []
    ys: List[int] = []
    for i in idx:
        x, y = ds[i]
        xs.append(x)
        ys.append(int(y))
    images = torch.stack(xs, dim=0).to(device)
    targets = torch.tensor(ys, dtype=torch.long, device=device)
    return to_neat_batch(images), to_binary_labels(targets, positive_class)


@torch.inference_mode()
def eval_planned_wave(
    items: Sequence[PopRec],
    images: torch.Tensor,
    labels: torch.Tensor,
    cache: ForwardActivationCache,
    *,
    activation: str,
) -> Tuple[Dict[GenomeId, float], Dict[str, Any], float]:
    """One planned unique-miss wave; return fitness + plan summary + wall_s."""
    device = images.device
    _sync(device)
    t0 = time.perf_counter()
    cache.precompute_population(list(items), register_overlay=True)
    work = cache.plan_population_misses(list(items))
    outputs = cache.execute_planned_population(
        images, work, store_activations=False, pin=False
    )
    _sync(device)
    wall = time.perf_counter() - t0

    fitness: Dict[GenomeId, float] = {}
    for genome, gid in items:
        out = outputs[gid].float()
        scores = output_to_score(out, output_readout="mean", activation=activation)
        preds = predict_from_scores(scores, decision_rule="score", decision_threshold=0.0)
        tp = ((labels == 1) & (preds == 1)).sum().item()
        tn = ((labels == 0) & (preds == 0)).sum().item()
        fp = ((labels == 0) & (preds == 1)).sum().item()
        fn = ((labels == 1) & (preds == 0)).sum().item()
        tpr = tp / max(1, tp + fn)
        tnr = tn / max(1, tn + fp)
        fitness[gid] = 0.5 * (tpr + tnr)
        _ = genome  # kept for type clarity / future use
    return fitness, work.summary(), wall


def _commit_elites(
    cache: ForwardActivationCache,
    elites: Sequence[PopRec],
    images: torch.Tensor,
) -> None:
    cache.precompute_population(list(elites), register_overlay=False)
    for genome, gid in elites:
        genome.eval()
        cache.execute(genome, images, gid, pin=True, store_activations=True)


def run_evolution(
    *,
    n_genomes: int,
    n_images: int,
    generations: int,
    activation: str,
    seed: int,
    device: torch.device,
    images: torch.Tensor,
    labels: torch.Tensor,
    max_cache_bytes: int,
    elite_percent: float = 0.2,
    verbose: bool = True,
    vram_limit_mb: float | None = None,
    max_gen_wall_s: float | None = None,
) -> RunResult:
    rng = random.Random(seed)
    torch.manual_seed(seed)
    if device.type == "cuda":
        torch.cuda.manual_seed_all(seed)

    result = RunResult(
        seed=seed,
        activation=activation,
        n_genomes=n_genomes,
        n_images=n_images,
        generations=generations,
        ok=False,
    )
    n_elite = max(1, int(round(n_genomes * elite_percent)))
    data_id = data_id_from_tensor(images)
    cache = ForwardActivationCache(max_bytes=max_cache_bytes, data_id=data_id)
    if vram_limit_mb is None:
        vram_limit_mb = 0.85 * _device_total_mb(device)

    try:
        _reset_peak(device)
        stem = _seed_genome(activation, rng)
        stem.to(device)
        stem.eval()
        elites = [stem.clone() for _ in range(n_elite)]
        for i, e in enumerate(elites):
            e = ConvolutionalNetworkMutator(e).apply(
                MutationType.MUTATE_WEIGHTS, rate=1.0, scale=0.02 + 0.01 * i
            )
            e.to(device)
            e.eval()
            elites[i] = e

        t_run = time.perf_counter()
        for gen in range(generations):
            pop = _breed(elites, pop_size=n_genomes, generation=gen, rng=rng)
            for g, _ in pop:
                g.to(device)
                g.eval()

            fitness, plan_sum, wall = eval_planned_wave(
                pop, images, labels, cache, activation=activation
            )
            peak_mb = _peak_vram_mb(device)
            if vram_limit_mb > 0 and peak_mb > vram_limit_mb:
                raise VramBudgetExceeded(
                    f"peak {peak_mb:.0f}MB > limit {vram_limit_mb:.0f}MB"
                )
            # Wall-time guard only when already using a large fraction of VRAM —
            # cold starts on small configs are slow but not capacity failures.
            if (
                max_gen_wall_s is not None
                and wall > max_gen_wall_s
                and peak_mb > 0.5 * vram_limit_mb
            ):
                raise VramBudgetExceeded(
                    f"gen wall {wall:.1f}s > max_gen_wall_s {max_gen_wall_s:.1f}s "
                    f"at peak_vram={peak_mb:.0f}MB (likely host spill)"
                )

            ranked = sorted(
                pop,
                key=lambda rec: (
                    fitness.get(rec[1], 0.0),
                    _count_edges(rec[0]),
                ),
                reverse=True,
            )
            elite_recs = ranked[:n_elite]
            elites = [g for g, _ in elite_recs]
            cache.set_watchlist({gid for _, gid in elite_recs})
            _commit_elites(cache, elite_recs, images)
            cache.end_generation()

            mean_nodes = sum(_count_nodes(g) for g, _ in pop) / max(1, len(pop))
            mean_edges = sum(_count_edges(g) for g, _ in pop) / max(1, len(pop))
            st = cache.stats
            gm = GenMetrics(
                generation=gen,
                n_genomes=n_genomes,
                n_images=n_images,
                mean_nodes=float(mean_nodes),
                mean_edges=float(mean_edges),
                hit_slots=int(plan_sum.get("hit_slots", 0)),
                miss_slots=int(plan_sum.get("miss_slots", 0)),
                unique_misses=int(plan_sum.get("unique_misses", 0)),
                dedup_saved=int(plan_sum.get("dedup_saved", 0)),
                hit_rate=float(st.hit_rate()),
                wall_s=float(wall),
                peak_vram_mb=float(peak_mb),
                cache_peak_mb=float(st.peak_bytes) / (1024.0 * 1024.0),
                cache_bytes_mb=float(st.bytes_used) / (1024.0 * 1024.0),
                elite_fitness=float(fitness.get(elite_recs[0][1], 0.0)),
            )
            result.gens.append(gm)
            if verbose and (gen % max(1, generations // 10) == 0 or gen + 1 == generations):
                print(
                    f"  gen {gen + 1:>3}/{generations}  "
                    f"nodes={gm.mean_nodes:.1f} edges={gm.mean_edges:.1f}  "
                    f"uniq_miss={gm.unique_misses} dedup={gm.dedup_saved}  "
                    f"hit={gm.hit_rate:.2f}  "
                    f"{gm.wall_s:.2f}s  vram={gm.peak_vram_mb:.0f}MB",
                    flush=True,
                )

        result.ok = True
        result.wall_s = time.perf_counter() - t_run
        result.peak_vram_mb = max((g.peak_vram_mb for g in result.gens), default=0.0)
        if result.gens:
            result.mean_unique_misses = sum(g.unique_misses for g in result.gens) / len(
                result.gens
            )
            result.mean_hit_rate = sum(g.hit_rate for g in result.gens) / len(result.gens)
            result.mean_wall_per_gen_s = sum(g.wall_s for g in result.gens) / len(
                result.gens
            )
            late = result.gens[-min(10, len(result.gens)) :]
            result.late_mean_nodes = sum(g.mean_nodes for g in late) / len(late)
            result.late_mean_edges = sum(g.mean_edges for g in late) / len(late)
    except (torch.cuda.OutOfMemoryError, VramBudgetExceeded) as exc:
        result.fail_reason = f"{type(exc).__name__}: {exc}"
        if device.type == "cuda":
            torch.cuda.empty_cache()
    except Exception as exc:  # noqa: BLE001 — capacity harness must not die
        result.fail_reason = f"{type(exc).__name__}: {exc}"
        if verbose:
            traceback.print_exc()
    return result


def probe_capacity(
    *,
    device: torch.device,
    cifar_root: Path,
    activation: str,
    max_cache_bytes: int,
    probe_gens: int,
    pop_grid: Sequence[int],
    image_grid: Sequence[int],
    seed: int,
    vram_limit_mb: float,
    max_gen_wall_s: float,
) -> Dict[str, Any]:
    """Geometric probe: max pop at fixed images, max images at fixed pop."""
    print("\n=== CAPACITY PROBE ===", flush=True)
    print(
        f"VRAM budget {vram_limit_mb:.0f}MB "
        f"(device {_device_total_mb(device):.0f}MB), "
        f"max_gen_wall_s={max_gen_wall_s:.1f}",
        flush=True,
    )
    # Mid-grid images for pop sweep — prefer a value that itself fits later.
    base_images = 512
    if image_grid:
        # Prefer 512 if present, else median.
        base_images = 512 if 512 in image_grid else int(image_grid[len(image_grid) // 2])
    base_pop = int(pop_grid[0]) if pop_grid else 32

    def _one(n_pop: int, n_img: int, *, tag: str, seed_i: int) -> CapacityPoint:
        print(
            f"\n[{tag}] genomes={n_pop} images={n_img} gens={probe_gens}",
            flush=True,
        )
        _reset_peak(device)
        try:
            images, labels = load_cifar_neat_batch(
                n_images=n_img, device=device, cifar_root=cifar_root, seed=seed_i
            )
        except torch.cuda.OutOfMemoryError as exc:
            if device.type == "cuda":
                torch.cuda.empty_cache()
            print("  FAIL OOM loading images", flush=True)
            return CapacityPoint(
                n_genomes=n_pop,
                n_images=n_img,
                ok=False,
                fail_reason=f"OOM loading images: {exc}",
            )
        rr = run_evolution(
            n_genomes=n_pop,
            n_images=n_img,
            generations=probe_gens,
            activation=activation,
            seed=seed_i,
            device=device,
            images=images,
            labels=labels,
            max_cache_bytes=max_cache_bytes,
            verbose=True,
            vram_limit_mb=vram_limit_mb,
            max_gen_wall_s=max_gen_wall_s,
        )
        del images, labels
        if device.type == "cuda":
            torch.cuda.empty_cache()
        pt = CapacityPoint(
            n_genomes=n_pop,
            n_images=n_img,
            ok=rr.ok,
            wall_s=rr.wall_s,
            peak_vram_mb=rr.peak_vram_mb,
            fail_reason=rr.fail_reason,
        )
        if rr.ok:
            print(f"  OK peak_vram={rr.peak_vram_mb:.0f}MB", flush=True)
        else:
            print(f"  FAIL {rr.fail_reason}", flush=True)
        return pt

    pop_points: List[CapacityPoint] = []
    best_pop = 0
    for n_pop in pop_grid:
        pt = _one(n_pop, base_images, tag="probe pop", seed_i=seed)
        pop_points.append(pt)
        if pt.ok:
            best_pop = n_pop
        else:
            break

    if best_pop <= 0:
        best_pop = max(8, base_pop // 2)

    # Image sweep at a moderate pop so we measure image capacity, not pop×images blowup.
    img_probe_pop = max(8, min(int(best_pop), 64))
    img_points: List[CapacityPoint] = []
    best_images = 0
    for n_img in image_grid:
        pt = _one(img_probe_pop, n_img, tag="probe img", seed_i=seed + 1)
        img_points.append(pt)
        if pt.ok:
            best_images = n_img
        else:
            break

    if best_images <= 0:
        best_images = max(64, base_images // 2)

    # Short probe underestimates long-run growth; keep ~50% for 100-gen confirms.
    confirm_pop = max(8, int(best_pop * 0.5))
    confirm_images = max(64, int(best_images * 0.5))
    confirm_pop = max(8, (confirm_pop // 8) * 8)
    confirm_images = max(64, (confirm_images // 64) * 64)

    summary = {
        "vram_limit_mb": vram_limit_mb,
        "max_gen_wall_s": max_gen_wall_s,
        "base_images_for_pop_sweep": base_images,
        "img_probe_pop": img_probe_pop,
        "best_pop": best_pop,
        "best_images": best_images,
        "confirm_pop": confirm_pop,
        "confirm_images": confirm_images,
        "pop_points": [asdict(p) for p in pop_points],
        "image_points": [asdict(p) for p in img_points],
        "note": (
            "best_* are short-probe (2 gens) frontiers on RTX onboard VRAM; "
            "confirm_* are ~50% for 100-gen topology growth headroom."
        ),
    }
    print(
        f"\nCapacity frontier: pop<={best_pop} @ images={base_images}; "
        f"images<={best_images} @ pop={img_probe_pop}",
        flush=True,
    )
    print(
        f"Confirm settings (~50% for 100-gen growth): "
        f"pop={confirm_pop} images={confirm_images}",
        flush=True,
    )
    return summary


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        description=(
            "Real CUDA planned-wave forward-cache capacity test "
            "(genomes x images), random/GELU activation, multi-run 100 gens."
        )
    )
    p.add_argument(
        "--activation",
        default="gelu",
        help="Activation: gelu|relu|tanh|leaky_relu|random (default gelu)",
    )
    p.add_argument("--runs", type=int, default=3, help="Confirm runs / seeds (default 3)")
    p.add_argument("--gens", type=int, default=100, help="Generations per run (default 100)")
    p.add_argument("--pop", type=int, default=0, help="Genomes/gen (0 = from probe)")
    p.add_argument("--images", type=int, default=0, help="Images/eval (0 = from probe)")
    p.add_argument("--probe-gens", type=int, default=2, help="Gens per probe step")
    p.add_argument("--no-probe", action="store_true", help="Skip capacity probe")
    p.add_argument("--probe-only", action="store_true", help="Only run capacity probe")
    p.add_argument(
        "--pop-grid",
        default="16,32,64,128,192,256",
        help="Comma list of population sizes to probe",
    )
    p.add_argument(
        "--image-grid",
        default="128,256,512,1024,1536,2048",
        help="Comma list of image counts to probe",
    )
    p.add_argument(
        "--cache-gb",
        type=float,
        default=8.0,
        help="ForwardActivationCache budget in GiB (default 8)",
    )
    p.add_argument(
        "--vram-frac",
        type=float,
        default=0.85,
        help="Treat peak allocated > frac*device_total as capacity failure (default 0.85)",
    )
    p.add_argument(
        "--max-gen-wall-s",
        type=float,
        default=20.0,
        help=(
            "Probe: fail a point if one gen exceeds this wall AND peak VRAM "
            "> 50% of budget (host-spill heuristic; default 20)"
        ),
    )
    p.add_argument("--seed", type=int, default=42)
    p.add_argument(
        "--cifar-root",
        type=Path,
        default=ROOT / "data" / "cifar10",
    )
    p.add_argument(
        "--out",
        type=Path,
        default=ROOT / "runs" / "cache_wave_sim" / "real_capacity",
    )
    return p


def _parse_int_grid(raw: str) -> List[int]:
    return [int(x.strip()) for x in raw.split(",") if x.strip()]


def main(argv: Optional[Sequence[str]] = None) -> int:
    args = build_parser().parse_args(argv)
    if not torch.cuda.is_available():
        print("CUDA required for real wave-cache capacity test.", file=sys.stderr)
        return 2
    device = torch.device("cuda")
    args.out.mkdir(parents=True, exist_ok=True)
    max_cache_bytes = int(args.cache_gb * 1024**3)
    vram_limit_mb = float(args.vram_frac) * _device_total_mb(device)
    max_gen_wall_s = float(args.max_gen_wall_s)
    # Hard allocator cap so oversized points fail fast with OOM instead of host spill.
    torch.cuda.set_per_process_memory_fraction(min(0.95, float(args.vram_frac) + 0.05))
    rng = random.Random(args.seed)
    activation0 = _pick_activation(args.activation, rng)

    report: Dict[str, Any] = {
        "device": torch.cuda.get_device_name(0),
        "vram_total_gb": round(torch.cuda.get_device_properties(0).total_memory / 1e9, 2),
        "vram_limit_mb": vram_limit_mb,
        "max_gen_wall_s": max_gen_wall_s,
        "activation_mode": args.activation,
        "cache_gb": args.cache_gb,
        "protocol": {
            "planned_wave": True,
            "cifar32": True,
            "mutations": [m.value for m in _MUTATION_BAG],
            "gens_per_run": args.gens,
            "runs": args.runs,
        },
    }

    pop = int(args.pop)
    images_n = int(args.images)
    capacity: Dict[str, Any] | None = None

    if not args.no_probe:
        capacity = probe_capacity(
            device=device,
            cifar_root=args.cifar_root,
            activation=activation0 if args.activation != "random" else "gelu",
            max_cache_bytes=max_cache_bytes,
            probe_gens=int(args.probe_gens),
            pop_grid=_parse_int_grid(args.pop_grid),
            image_grid=_parse_int_grid(args.image_grid),
            seed=int(args.seed),
            vram_limit_mb=vram_limit_mb,
            max_gen_wall_s=max_gen_wall_s,
        )
        report["capacity"] = capacity
        if pop <= 0:
            pop = int(capacity["confirm_pop"])
        if images_n <= 0:
            images_n = int(capacity["confirm_images"])
    else:
        if pop <= 0:
            pop = 64
        if images_n <= 0:
            images_n = 512

    if args.probe_only:
        out_path = args.out / "capacity_report.json"
        out_path.write_text(json.dumps(report, indent=2), encoding="utf-8")
        print(f"\nWrote {out_path}")
        return 0

    print(
        f"\n=== CONFIRM: {args.runs} runs x {args.gens} gens | "
        f"pop={pop} images={images_n} activation={args.activation} ===",
        flush=True,
    )
    runs: List[Dict[str, Any]] = []
    for i in range(int(args.runs)):
        seed = int(args.seed) + 1000 * (i + 1)
        act = _pick_activation(args.activation, random.Random(seed))
        print(
            f"\n--- run {i + 1}/{args.runs} seed={seed} activation={act} ---",
            flush=True,
        )
        imgs, labs = load_cifar_neat_batch(
            n_images=images_n,
            device=device,
            cifar_root=args.cifar_root,
            seed=seed,
        )
        rr = run_evolution(
            n_genomes=pop,
            n_images=images_n,
            generations=int(args.gens),
            activation=act,
            seed=seed,
            device=device,
            images=imgs,
            labels=labs,
            max_cache_bytes=max_cache_bytes,
            verbose=True,
            vram_limit_mb=vram_limit_mb,
            max_gen_wall_s=None,  # confirm: allow slower late gens, still VRAM-capped
        )
        del imgs, labs
        payload = asdict(rr)
        runs.append(payload)
        status = "OK" if rr.ok else f"FAIL ({rr.fail_reason})"
        print(
            f"  => {status} wall={rr.wall_s:.1f}s peak_vram={rr.peak_vram_mb:.0f}MB "
            f"mean_hit={rr.mean_hit_rate:.2f} late_edges={rr.late_mean_edges:.1f}",
            flush=True,
        )
        if device.type == "cuda":
            torch.cuda.empty_cache()

    ok_runs = [r for r in runs if r["ok"]]
    report["confirm"] = {
        "pop": pop,
        "images": images_n,
        "runs": runs,
        "n_ok": len(ok_runs),
        "n_fail": len(runs) - len(ok_runs),
    }
    if ok_runs:
        report["confirm"]["aggregate"] = {
            "mean_wall_s": sum(r["wall_s"] for r in ok_runs) / len(ok_runs),
            "mean_peak_vram_mb": sum(r["peak_vram_mb"] for r in ok_runs) / len(ok_runs),
            "mean_hit_rate": sum(r["mean_hit_rate"] for r in ok_runs) / len(ok_runs),
            "mean_unique_misses": sum(r["mean_unique_misses"] for r in ok_runs)
            / len(ok_runs),
            "mean_wall_per_gen_s": sum(r["mean_wall_per_gen_s"] for r in ok_runs)
            / len(ok_runs),
            "late_mean_nodes": sum(r["late_mean_nodes"] for r in ok_runs) / len(ok_runs),
            "late_mean_edges": sum(r["late_mean_edges"] for r in ok_runs) / len(ok_runs),
            "genomes_per_generation": pop,
            "images_per_eval": images_n,
            "images_x_genomes_per_gen": pop * images_n,
        }

    out_path = args.out / "real_capacity_report.json"
    out_path.write_text(json.dumps(report, indent=2), encoding="utf-8")
    print(f"\nWrote {out_path}", flush=True)

    print("\n=== SUMMARY ===", flush=True)
    if capacity:
        print(
            f"Capacity frontier: <= {capacity['best_pop']} genomes @ "
            f"{capacity['base_images_for_pop_sweep']} imgs; "
            f"<= {capacity['best_images']} images @ {capacity['best_pop']} genomes",
            flush=True,
        )
    print(f"Confirm: {pop} genomes x {images_n} images, {args.gens} gens", flush=True)
    print(f"Runs OK: {len(ok_runs)}/{len(runs)}", flush=True)
    if ok_runs and "aggregate" in report["confirm"]:
        agg = report["confirm"]["aggregate"]
        print(
            f"Agg: {agg['mean_wall_per_gen_s']:.2f}s/gen | "
            f"hit={agg['mean_hit_rate']:.2f} | "
            f"vram={agg['mean_peak_vram_mb']:.0f}MB | "
            f"throughput~{agg['images_x_genomes_per_gen']} image-forwards/gen",
            flush=True,
        )
    return 0 if ok_runs else 1


if __name__ == "__main__":
    raise SystemExit(main())