tools/build_all.py

"""
Build and verify every named (bits, memory_profile) variant.

Outputs:
  variants/neural_alu{8,16,32}.safetensors                          - no memory
  variants/neural_computer{8,16,32}_registers.safetensors           - 16 B
  variants/neural_computer{8,16,32}_scratchpad.safetensors          - 256 B
  variants/neural_computer{8,16,32}_small.safetensors               - 1 KB
  variants/neural_computer{8,16,32}_reduced.safetensors             - 4 KB
  variants/neural_computer{8,16,32}.safetensors                     - 64 KB

Each variant is built from the canonical seed, quantized to minimum
integer dtypes with the strict ternary check (--ternary --strict), and
then verified with eval.py via the BatchedFitnessEvaluator; the summary
records (tensor count, params, file size, fitness, total_tests, seconds).
"""

from __future__ import annotations
import os
import shutil
import subprocess
import sys
import time
from pathlib import Path

import torch
from safetensors import safe_open

ROOT = Path(__file__).resolve().parent.parent  # repo root (this file lives in tools/)
SEED = ROOT / "neural_computer.safetensors"
OUT_DIR = ROOT / "variants"
OUT_DIR.mkdir(exist_ok=True)

PROFILES = ["none", "registers", "scratchpad", "small", "reduced", "full"]
BITS = [8, 16, 32]


def variant_filename(bits: int, profile: str) -> str:
    if profile == "none":
        return f"neural_alu{bits}.safetensors"
    if profile == "full":
        return f"neural_computer{bits}.safetensors"
    return f"neural_computer{bits}_{profile}.safetensors"


def run(cmd: list[str], timeout: int = 600) -> tuple[int, str]:
    p = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
    return p.returncode, (p.stdout or "") + (p.stderr or "")


def build_variant(bits: int, profile: str) -> Path:
    out = OUT_DIR / variant_filename(bits, profile)
    shutil.copy2(SEED, out)
    cmd = [
        sys.executable, str(ROOT / "src" / "build.py"),
        "--bits", str(bits),
        "-m", profile,
        "--apply",
        "--model", str(out),
        "all",
    ]
    rc, log = run(cmd, timeout=2400)
    if rc != 0:
        raise RuntimeError(f"build failed for bits={bits} profile={profile}:\n{log[-1500:]}")
    return out


def quantize_variant(path: Path) -> tuple[int, int]:
    """Run quantize.py --ternary --strict on a built variant. This is the
    last pipeline step: it casts every tensor to its minimum signed integer
    dtype, verifies the strictly ternary weight invariant, and stamps the
    weight_quantization metadata field. Returns (bytes_before, bytes_after).
    """
    rc, log = run([sys.executable, str(ROOT / "src" / "quantize.py"), str(path),
                   "--ternary", "--strict"], timeout=300)
    if rc != 0:
        raise RuntimeError(f"quantize failed for {path.name}:\n{log[-800:]}")
    # parse the "file X.X MB -> Y.Y MB" line
    for line in log.splitlines():
        if "file" in line and "->" in line and path.name in line:
            try:
                parts = line.split("file")[1].split("->")
                before = float(parts[0].strip().split()[0]) * 1e6
                after = float(parts[1].strip().split()[0]) * 1e6
                return int(before), int(after)
            except Exception:
                pass
    return 0, 0


def measure_variant(path: Path) -> dict:
    """Read tensor count, params, manifest values from the variant."""
    with safe_open(str(path), framework="pt") as f:
        keys = list(f.keys())
        params = sum(f.get_tensor(k).numel() for k in keys)
        manifest = {
            k.split(".", 1)[1]: f.get_tensor(k).item()
            for k in keys if k.startswith("manifest.") and f.get_tensor(k).numel() == 1
        }
    return {
        "tensors": len(keys),
        "params": params,
        "size_mb": path.stat().st_size / (1024 * 1024),
        "manifest": manifest,
    }


def eval_variant(path: Path, device: str = "cpu", timeout: int = 600) -> dict:
    """Run eval.py against a variant and parse fitness."""
    cmd = [
        sys.executable, str(ROOT / "src" / "eval.py"),
        "--model", str(path),
        "--device", device,
        "--quiet",
    ]
    t0 = time.time()
    rc, log = run(cmd, timeout=timeout)
    dt = time.time() - t0

    fitness = None
    total_tests = None
    status = "ERROR"
    for line in log.splitlines():
        line = line.strip()
        if line.startswith("Fitness:"):
            try:
                fitness = float(line.split()[1])
            except Exception:
                pass
        elif line.startswith("Total tests:"):
            try:
                total_tests = int(line.split()[2])
            except Exception:
                pass
        elif line.startswith("STATUS:"):
            status = line.split()[1]
    return {
        "rc": rc,
        "fitness": fitness,
        "total_tests": total_tests,
        "status": status,
        "elapsed_s": dt,
        "log_tail": "\n".join(log.splitlines()[-15:]),
    }


def main() -> None:
    rows = []
    print(f"Building 18 variants into {OUT_DIR}\n")
    for bits in BITS:
        for profile in PROFILES:
            label = f"bits={bits} profile={profile}"
            print(f"=== {label} ===", flush=True)
            t0 = time.time()
            try:
                path = build_variant(bits, profile)
                bt = time.time() - t0
                pre_q_meta = measure_variant(path)
                # Quantize in-place as the final step; weights are
                # integer-valued so this is exact, --strict fails the build
                # if any weight is non-ternary, and header metadata
                # (signal_registry) is carried through.
                qb, qa = quantize_variant(path)
                meta = measure_variant(path)
                ev = eval_variant(path, device="cpu", timeout=900)
                rows.append({
                    "bits": bits, "profile": profile,
                    "filename": path.name,
                    "build_s": bt,
                    **meta,
                    **{k: ev[k] for k in ("fitness", "total_tests", "status", "elapsed_s")},
                    "log_tail": ev["log_tail"] if ev["status"] != "PASS" else "",
                })
                q_ratio = qb / qa if qa else 1.0
                print(f"  built in {bt:.1f}s  size={pre_q_meta['size_mb']:.1f}MB -> "
                      f"{meta['size_mb']:.1f}MB after quant ({q_ratio:.2f}x)"
                      f"  params={meta['params']:,}  tensors={meta['tensors']:,}")
                print(f"  eval: fitness={ev['fitness']}  tests={ev['total_tests']}"
                      f"  status={ev['status']}  ({ev['elapsed_s']:.1f}s)")
                if ev["status"] != "PASS":
                    print("  --- failure tail ---")
                    print("  " + "\n  ".join(ev["log_tail"].splitlines()))
                    print("  --------------------")
            except Exception as e:
                print(f"  EXCEPTION: {e}")
                rows.append({"bits": bits, "profile": profile, "error": str(e)})
            print()

    print("=" * 88)
    print(" SUMMARY")
    print("=" * 88)
    header = f"{'bits':>4} {'profile':<11} {'size_MB':>8} {'tensors':>8} {'params':>11} {'fitness':>9} {'tests':>6} {'status':>7}"
    print(header)
    print("-" * len(header))
    for r in rows:
        if "error" in r:
            print(f"{r['bits']:>4} {r['profile']:<11} ERROR: {r['error'][:60]}")
            continue
        fit = f"{r['fitness']:.4f}" if r['fitness'] is not None else "n/a"
        tests = r['total_tests'] if r['total_tests'] is not None else "?"
        print(f"{r['bits']:>4} {r['profile']:<11} {r['size_mb']:>8.1f} "
              f"{r['tensors']:>8,} {r['params']:>11,} "
              f"{fit:>9} {tests:>6} {r['status']:>7}")

    fail = [r for r in rows if r.get("status") != "PASS" or "error" in r]
    print()
    if fail:
        print(f"FAILURES: {len(fail)}/{len(rows)}")
    else:
        print(f"ALL {len(rows)} VARIANTS PASS")


if __name__ == "__main__":
    main()