digital-drosophila/models/male-cns-lif/runtime-credit-v1/phase4o_a.py

#!/usr/bin/env python3
"""Read-only Phase 4O-A runtime determinism audit; never writes frozen roots."""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import platform
import random
import shutil
import subprocess
import sys
import time
from pathlib import Path

import numpy as np
import pyarrow.parquet as pq
import torch
import yaml

ROOT = Path("/data/digital-drosophila")
EXP = ROOT / "experiments/phase4o-runtime-credit"
MODEL = ROOT / "models/male-cns-lif/runtime-credit-v1"
FROZEN = ROOT / "experiments/phase4c-runtime/regression/reference"
GPU_UUID = "GPU-79f3a85d-adce-7833-7003-6488140ca596"


def save(path: Path, value: object) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True, default=str) + "\n")


def sha(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        for block in iter(lambda: f.read(8 << 20), b""):
            h.update(block)
    return h.hexdigest()


def digest(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def cmd(args: list[str]) -> dict:
    exe = shutil.which(args[0])
    if exe is None:
        return {"available": False}
    p = subprocess.run(args, text=True, capture_output=True)
    return {"available": True, "exit_code": p.returncode, "stdout": p.stdout.strip()[:8000], "stderr": p.stderr.strip()[:2000]}


def env_audit() -> None:
    if os.environ.get("CUDA_VISIBLE_DEVICES") != GPU_UUID:
        raise RuntimeError("CUDA_VISIBLE_DEVICES must select frozen RTX 4090 UUID")
    names = ["CUDA_VISIBLE_DEVICES", "CUBLAS_WORKSPACE_CONFIG", "PYTHONHASHSEED", "OMP_NUM_THREADS", "MKL_NUM_THREADS"]
    names += sorted(x for x in os.environ if x.startswith("TORCH_"))
    gpu = torch.cuda.get_device_properties(0)
    e = {"OS_release": platform.platform(), "kernel": platform.release(), "uname": platform.uname()._asdict(),
         "python": sys.version, "torch": torch.__version__, "cuda_runtime": torch.version.cuda,
         "cudnn": torch.backends.cudnn.version(), "gpu_name": gpu.name, "gpu_uuid": GPU_UUID,
         "gpu_total_memory_bytes": gpu.total_memory, "env_allowlist": {x: os.environ.get(x) for x in names},
         "nvcc": cmd(["nvcc", "--version"]), "nvidia_smi": cmd(["nvidia-smi", "--query-gpu=name,uuid,driver_version,memory.total", "--format=csv,noheader"]),
         "gcc": cmd(["gcc", "--version"]), "gxx": cmd(["g++", "--version"]), "ninja": cmd(["ninja", "--version"]),
         "git_model": cmd(["git", "-C", str(ROOT / "models/male-cns-lif"), "status", "--short", "--branch"]),
         "git_workspace": cmd(["git", "-C", str(ROOT), "rev-parse", "HEAD"]),
         "torch_flags": {"deterministic_algorithms": torch.are_deterministic_algorithms_enabled(),
                         "cudnn_deterministic": torch.backends.cudnn.deterministic,
                         "cudnn_benchmark": torch.backends.cudnn.benchmark,
                         "cuda_matmul_allow_tf32": torch.backends.cuda.matmul.allow_tf32,
                         "cudnn_allow_tf32": torch.backends.cudnn.allow_tf32,
                         "matmul_precision": torch.get_float32_matmul_precision()}}
    save(EXP / "environment/runtime-environment.json", e)
    roots = [ROOT / "models/male-cns-lif/reference-v1", ROOT / "models/male-cns-lif/runtime-v1",
             FROZEN, ROOT / "models/male-cns-lif/kcmbon-efficacy-v1",
             ROOT / "experiments/phase4n-credit-assignment/compiler-reports/extension-build",
             ROOT / "experiments/phase4c-runtime/compiler-reports/extension-build"]
    files = []
    for root in roots:
        if root.exists():
            for p in sorted(root.rglob("*")):
                if p.is_file() and p.suffix not in (".pyc", ".o"):
                    files.append({"path": str(p), "bytes": p.stat().st_size, "mtime_ns": p.stat().st_mtime_ns, "sha256": sha(p)})
    for p in [ROOT / "models/male-cns-lif/brain_runtime.py", ROOT / "models/male-cns-lif/phase4i.py",
              ROOT / "models/male-cns-lif/phase4m.py", ROOT / "models/male-cns-lif/phase4l.py",
              ROOT / "models/male-cns-lif/phase4n_events.py"]:
        if p.is_file():
            files.append({"path": str(p), "bytes": p.stat().st_size, "mtime_ns": p.stat().st_mtime_ns, "sha256": sha(p)})
    save(EXP / "binary-audit/execution-file-sha256.json", files)
    binary = [x for x in files if x["path"].endswith(".so")]
    save(EXP / "binary-audit/runtime-binary-diff-report.json", {
        "binary_files": binary, "bytewise_same": len({x["sha256"] for x in binary}) == 1 if len(binary) > 1 else None,
        "warning": "A binary hash difference is not itself proof of changed semantics; build paths and toolchain can alter ELF bytes.",
        "historical_phase4c_manifest": str(ROOT / "experiments/phase4c-runtime/manifests/manifest.json")})


def rng_hashes(rt=None) -> dict:
    py = digest(repr(random.getstate()).encode())
    npy = digest(repr(np.random.get_state()).encode())
    out = {"python_random": py, "numpy_global": npy, "torch_cpu": digest(torch.get_rng_state().numpy().tobytes()),
           "torch_cuda": digest(torch.cuda.get_rng_state(0).cpu().numpy().tobytes())}
    if rt is not None and hasattr(rt, "generator"):
        out["input_generator"] = digest(rt.generator.get_state().cpu().numpy().tobytes())
    return out


def first_event_diff(a_step, a_node, b_step, b_node):
    n = min(len(a_step), len(b_step))
    mismatched = np.flatnonzero((a_step[:n] != b_step[:n]) | (a_node[:n] != b_node[:n]))
    j = int(mismatched[0]) if len(mismatched) else (n if len(a_step) != len(b_step) else None)
    if j is None:
        return None
    ref = None if j >= len(a_step) else {"step": int(a_step[j]), "neuron_idx": int(a_node[j])}
    got = None if j >= len(b_step) else {"step": int(b_step[j]), "neuron_idx": int(b_node[j])}
    step = min(x["step"] for x in (ref, got) if x is not None)
    return {"event_position": j, "first_difference_step_upper_bound": step, "reference_event": ref, "observed_event": got,
            "note": "Event-list first difference is not necessarily first state divergence; state-level instrumentation is required."}


def repeat_audit(repeats: int, trial_filter: str | None = None, tag: str = "baseline") -> None:
    from brain_runtime import BrainRuntime
    from phase4i import poisson_batch
    cfgpath = ROOT / "models/male-cns-lif/mb-mechanistic-v1/phase4i.yaml"
    rt = BrainRuntime(cfgpath, EXP / "compiler-reports/extension-build")
    backend_file = Path(rt.extension.__file__)
    identity = {"module": rt.extension.__name__, "binary_path": str(backend_file), "binary_sha256": sha(backend_file),
                "runtime_v1_python_sha256": sha(ROOT / "models/male-cns-lif/runtime-v1/phase4c.py"),
                "source": "BrainRuntime._hybrid_edges plus GPU scatter_add_", "config": str(cfgpath), "config_sha256": sha(cfgpath)}
    save(EXP / "binary-audit/loaded-backend.json", identity)
    specs = [(s["name"], s["population"], float(s["rate_hz"]), int(s["seed"])) for s in yaml.safe_load((ROOT / "models/male-cns-lif/runtime-v1/phase4c.yaml").read_text())["reference_trials"]]
    if trial_filter:
        specs = [x for x in specs if x[0] == trial_filter]
        if not specs: raise ValueError(f"unknown frozen trial {trial_filter}")
    rows = []
    output_dir = EXP / "phase4c-replay" / tag
    for name, pop, rate, seed in specs:
        refroot = FROZEN / name
        table = pq.read_table(refroot / "spike-events.parquet")
        rstep = table["step"].to_numpy().astype(np.int32); rnode = table["neuron_idx"].to_numpy().astype(np.int32)
        rcount = np.load(refroot / "per-neuron-spike-counts.int32.npy")
        rv = np.load(refroot / "final-voltage.float32.npy"); rg = np.load(refroot / "final-g.float32.npy")
        targets = np.flatnonzero(rt.population_masks[pop]).astype(np.int64)
        for repeat in range(repeats):
            before = rng_hashes(rt)
            rt.reset(seed)
            after_reset = rng_hashes(rt)
            before_first = rng_hashes(rt)
            t0 = time.perf_counter()
            error = None
            try:
                rt.step(poisson_batch(targets, rate, 1000.0, name), 1000.0)
                events = rt.read_spikes()
                step = events["step"].astype(np.int32); node = events["neuron_idx"].astype(np.int32)
                counts = rt.spike_counts.cpu().numpy().astype(np.int32)
                voltage = rt.v.cpu().numpy(); g = rt.g.cpu().numpy()
                match_events = bool(np.array_equal(step, rstep) and np.array_equal(node, rnode))
                match_counts = bool(np.array_equal(counts, rcount))
                row = {"trial": name, "repeat": repeat, "target_count": int(len(targets)),
                       "spike_events_exact": match_events, "per_neuron_counts_exact": match_counts,
                       "final_voltage_max_abs_error": float(np.max(np.abs(voltage-rv))),
                       "final_g_max_abs_error": float(np.max(np.abs(g-rg))),
                       "event_count": len(node), "event_sha256": digest(step.tobytes()+node.tobytes()),
                       "counts_sha256": digest(counts.tobytes()), "first_event_difference": first_event_diff(rstep,rnode,step,node),
                       "wall_seconds": time.perf_counter()-t0,
                       "rng_before_reset": before, "rng_after_reset": after_reset,
                       "rng_before_first_timestep": before_first, "rng_after_trial": rng_hashes(rt)}
                if not match_events:
                    output_dir.mkdir(parents=True, exist_ok=True)
                    pq.write_table(__import__('pyarrow').table({'step':step,'neuron_idx':node}),
                                   output_dir / f"divergent-{name}-repeat-{repeat}-events.parquet",compression='zstd')
                    np.savez_compressed(output_dir / f"divergent-{name}-repeat-{repeat}-state.npz",
                                        final_voltage=voltage, final_g=g, counts=counts)
            except Exception as exc:
                error = repr(exc)
                row = {"trial": name, "repeat": repeat, "error": error, "wall_seconds": time.perf_counter()-t0,
                       "rng_before_reset": before, "rng_after_reset": after_reset}
            rows.append(row)
            save(output_dir / "reproducibility-matrix.json", rows)
            print(json.dumps({"trial": name, "repeat": repeat, "event_exact": row.get("spike_events_exact"),
                              "count_exact": row.get("per_neuron_counts_exact"), "first_diff": row.get("first_event_difference"), "error": error}), flush=True)
    save(output_dir / "summary.json", {"repeats_per_trial": repeats, "all_spike_events_exact": all(r.get("spike_events_exact",False) for r in rows),
         "all_counts_exact": all(r.get("per_neuron_counts_exact",False) for r in rows),
         "trial_summary": {name: {"event_exact": sum(r.get("spike_events_exact",False) for r in rows if r["trial"]==name),
                                  "count_exact": sum(r.get("per_neuron_counts_exact",False) for r in rows if r["trial"]==name)} for name,_,_,_ in specs},
         "backend": identity})


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("task", choices=["environment", "repeats"])
    ap.add_argument("--repeats", type=int, default=10)
    ap.add_argument("--trial")
    ap.add_argument("--tag", default="baseline")
    args = ap.parse_args()
    EXP.mkdir(parents=True, exist_ok=True); MODEL.mkdir(parents=True, exist_ok=True)
    if args.task == "environment": env_audit()
    else: repeat_audit(args.repeats, args.trial, args.tag)


if __name__ == "__main__":
    main()