path: root/experiments/analyze_resnet_crossover_r2.py

#!/usr/bin/env python3
"""Audit the complete failure-retaining 27-cell ResNet R2 panel."""
import argparse
import hashlib
import json
import math
import os
import sys


ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from experiments.resnet_crossover_r2 import (
    DEPTHS,
    METHODS,
    DEFAULT_SELECTOR,
    r2_jobs,
    registry_sha256,
    selector_report,
)


def sha256(path):
    digest = hashlib.sha256()
    with open(path, "rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def read_json(path):
    with open(path, encoding="utf-8") as handle:
        return json.load(handle)


def finite(value):
    return (
        isinstance(value, (int, float))
        and not isinstance(value, bool)
        and math.isfinite(float(value))
    )


def expected_mode(method):
    return {
        "bp": "bp",
        "fa": "hfa",
        "dfa": "dfa",
        "clean_kp": "kp",
        "sdil": "kp_traffic",
    }[method]


def history_metrics(record, job):
    method = job["method"]
    expected_epochs = job["epochs"]
    losses = []
    validation_accuracies = []
    history_examples = 0
    if method == "ff":
        layers = record.get("layers", [])
        layer_indices = [row.get("layer") for row in layers]
        assert layer_indices == list(range(len(layers)))
        for layer in layers:
            epochs = layer.get("epochs", [])
            assert [row.get("epoch") for row in epochs] == list(
                range(1, len(epochs) + 1)
            )
            assert len(epochs) <= expected_epochs
            losses.extend(row.get("loss") for row in epochs)
            history_examples += sum(
                int(row.get("examples", 0)) for row in epochs
            )
        complete = (
            len(layers) == job["depth"]
            and all(
                len(layer.get("epochs", [])) == expected_epochs
                for layer in layers
            )
        )
        completed_epochs = sum(
            len(layer.get("epochs", [])) for layer in layers
        )
    else:
        epochs = record.get("epochs", [])
        assert [row.get("epoch") for row in epochs] == list(
            range(1, len(epochs) + 1)
        )
        assert len(epochs) <= expected_epochs
        complete = len(epochs) == expected_epochs
        completed_epochs = len(epochs)
        if method in ("pepita", "ep", "dualprop"):
            losses = [row.get("train_loss") for row in epochs]
            history_examples = sum(
                int(row.get("train_examples", 0)) for row in epochs
            )
            validation_accuracies = [
                row["validation"]["accuracy"]
                for row in epochs
                if row.get("validation") is not None
            ]
        else:
            losses = [row.get("train_loss") for row in epochs]
            history_examples = sum(
                int(row.get("train_examples", 0)) for row in epochs
            )
            validation_accuracies = [
                row["eval_accuracy"]
                for row in epochs
                if row.get("eval_accuracy") is not None
            ]
    final = record.get("final", {})
    final_accuracy = final.get("accuracy")
    final_loss = final.get("loss")
    finite_accuracies = [
        float(value)
        for value in validation_accuracies + [final_accuracy]
        if finite(value)
    ]
    best_accuracy = max(finite_accuracies) if finite_accuracies else None
    losses_finite = bool(losses) and all(finite(value) for value in losses)
    final_finite = (
        final.get("finite") is True
        and finite(final_accuracy)
        and (final_loss is None or finite(final_loss))
    )
    first_nonfinite = record.get("first_nonfinite_step")
    return {
        "complete_trajectory": complete,
        "epochs_completed": completed_epochs,
        "history_training_examples": history_examples,
        "all_training_losses_finite": losses_finite,
        "final_validation_accuracy":
            float(final_accuracy) if finite(final_accuracy) else None,
        "best_validation_accuracy":
            float(best_accuracy) if best_accuracy is not None else None,
        "final_validation_loss":
            float(final_loss) if finite(final_loss) else None,
        "first_nonfinite_step": first_nonfinite,
        "finite": (
            complete
            and losses_finite
            and final_finite
            and first_nonfinite is None
        ),
    }


def work_metrics(record, job, history):
    method = job["method"]
    if method in ("pepita", "ff", "ep", "dualprop"):
        work = record["work"]
        ordinary = int(work["ordinary_training_examples"])
        validation = int(work["ordinary_validation_examples"])
        presentations = int(work["training_example_presentations"])
        relaxation = int(work["relaxation_example_passes"])
        candidate = int(work["candidate_label_evaluation_presentations"])
        queries = int(work["logical_task_loss_queries"])
        local_vjp = int(work["local_vjp_example_evaluations"])
        local_backward = int(
            work["local_target_backward_example_evaluations"]
        )
        assert ordinary == history["history_training_examples"]
        if method == "ff":
            assert presentations == 2 * ordinary
            assert candidate == 10 * validation
            assert relaxation == 0
            assert local_backward == 2 * ordinary
            assert local_vjp == 0
        elif method == "pepita":
            assert presentations == 2 * ordinary
            assert relaxation == candidate == local_vjp == local_backward == 0
        elif method == "ep":
            assert presentations == ordinary
            assert relaxation == 24 * ordinary + 20 * validation
            assert local_vjp == relaxation
            assert candidate == local_backward == 0
        else:
            assert presentations == ordinary
            assert relaxation == 16 * ordinary
            assert local_vjp == relaxation
            assert candidate == local_backward == 0
        total_work = {
            "forward_macs_per_example": work["forward_macs_per_example"],
            "feedforward_example_passes": work["feedforward_example_passes"],
            "relaxation_example_passes": relaxation,
            "local_vjp_example_evaluations": local_vjp,
            "candidate_label_evaluation_presentations": candidate,
        }
    else:
        work = record["work"]
        counters = record["counters"]
        ordinary = int(counters["ordinary_examples"])
        validation = (
            int(record["split"]["validation_examples"])
            * int(record["evaluation_protocol"]["validation_evaluations"])
        )
        presentations = ordinary
        queries = int(work["logical_batch_loss_queries"])
        total_work = {
            "forward_macs_per_example": work["forward_macs_per_example"],
            "total_macs_estimate": work["total_macs_estimate"],
            "elementwise_operations_estimate":
                work["elementwise_operations_estimate"],
            "total_forward_equivalent_examples":
                work["total_forward_equivalent_examples"],
        }
        assert ordinary == history["history_training_examples"]
    assert queries == 0
    return {
        "ordinary_training_examples": ordinary,
        "ordinary_validation_examples": validation,
        "training_example_presentations": presentations,
        "logical_task_loss_queries": queries,
        "work": total_work,
    }


def audit_job(job, expected_source, expected_selector):
    manifest_path = job["output"] + ".manifest.json"
    if not os.path.isfile(manifest_path):
        raise AssertionError(f"missing R2 manifest: {job['experiment_name']}")
    manifest = read_json(manifest_path)
    for key in (
        "stage",
        "method",
        "architecture",
        "depth",
        "rate",
        "epochs",
        "lr_schedule",
        "experiment_name",
        "output",
        "timeout_seconds",
        "command",
    ):
        assert manifest[key] == job[key], (
            f"{job['experiment_name']}: manifest drift in {key}"
        )
    assert manifest["source"] == expected_source
    assert manifest["selector"] == expected_selector
    hardware = manifest["hardware_lock"]
    assert hardware["physical_gpu_index"] in (5, 7)
    assert hardware["physical_gpu_uuid"]
    common = {
        "cell_id": f"resnet{job['depth']}::{job['method']}",
        "method": job["method"],
        "depth": job["depth"],
        "rate": job["rate"],
        "status": manifest["status"],
        "manifest": os.path.relpath(manifest_path, ROOT),
        "driver_wall_seconds": float(manifest["driver_wall_seconds"]),
        "physical_gpu_index": hardware["physical_gpu_index"],
        "physical_gpu_uuid": hardware["physical_gpu_uuid"],
        "output_sha256": manifest["output_sha256"],
    }
    if manifest["status"] != "completed":
        assert manifest["status"] in {
            "timeout", "nonzero_exit", "missing_output"
        }
        assert (
            manifest["output_sha256"] is None
            if not manifest["output_exists"]
            else manifest["output_sha256"] == sha256(job["output"])
        )
        return {
            **common,
            "finite": False,
            "complete_trajectory": False,
            "epochs_completed": 0,
            "best_validation_accuracy": None,
            "final_validation_accuracy": None,
            "final_validation_loss": None,
            "first_nonfinite_step": None,
            "ordinary_training_examples": None,
            "ordinary_validation_examples": None,
            "training_example_presentations": None,
            "logical_task_loss_queries": None,
            "forward_parameter_count": None,
            "peak_memory_allocated_bytes": None,
            "total_wall_seconds": None,
            "work": None,
        }
    assert manifest["output_exists"] is True
    assert manifest["output_sha256"] == sha256(job["output"])
    record = read_json(job["output"])
    provenance = record["provenance"]
    assert provenance["git_commit"] == expected_source["git_commit"]
    assert provenance["git_tracked_dirty"] is False
    args = record["args"]
    for key, expected in {
        "depth": job["depth"],
        "width": 16,
        "seed": 0,
        "loader_seed": 0,
        "split_seed": 2027,
        "batch_size": 128,
        "epochs": job["epochs"],
        "train_limit": 0,
        "val_examples": 5000,
        "eval_split": "validation",
        "eval_every": 1,
        "augment_train": 1,
        "lr": job["rate"],
        "lr_schedule": job["lr_schedule"],
        "momentum": 0.9,
        "weight_decay": 1e-4,
    }.items():
        assert args[key] == expected, (
            f"{job['experiment_name']}: argument drift in {key}"
        )
    if job["method"] in ("pepita", "ff", "ep", "dualprop"):
        assert args["method"] == job["method"]
    else:
        assert args["mode"] == expected_mode(job["method"])
    split = record["split"]
    assert split["train_examples"] == 45_000
    assert split["validation_examples"] == 5_000
    assert split["split_seed"] == 2027
    evaluation = record["evaluation_protocol"]
    assert evaluation["test_evaluations"] == 0
    assert evaluation["test_used_for_selection"] is False
    architecture = record["architecture"]
    assert architecture["depth"] == job["depth"]
    assert architecture["base_width"] == 16
    history = history_metrics(record, job)
    work = work_metrics(record, job, history)
    if history["complete_trajectory"]:
        expected_ordinary = 45_000 * job["epochs"]
        if job["method"] == "ff":
            expected_ordinary *= job["depth"]
        assert work["ordinary_training_examples"] == expected_ordinary
    hardware_record = record["hardware"]
    assert hardware_record["cuda_visible_devices"] in ("5", "7")
    assert hardware_record["cuda_device_name"] == "NVIDIA GeForce GTX 1080"
    parameter_count = architecture.get("forward_parameter_count")
    if parameter_count is None:
        parameter_count = architecture["forward_parameters"]
    return {
        **common,
        **history,
        **work,
        "forward_parameter_count": int(parameter_count),
        "peak_memory_allocated_bytes": int(
            hardware_record["peak_memory_allocated_bytes"]
        ),
        "total_wall_seconds": float(
            record.get("total_wall_seconds")
            if "total_wall_seconds" in record
            else record["timing"]["total_timed_wall_s"]
        ),
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--selector", default=DEFAULT_SELECTOR)
    parser.add_argument(
        "--out", default="results/resnet_crossover/r2_audit.json"
    )
    args = parser.parse_args()
    selector = selector_report(args.selector)
    jobs = r2_jobs(selector["selected_rates"])
    launch_path = os.path.join(
        ROOT, "results", "resnet_crossover", "r2_launch.json"
    )
    assert os.path.isfile(launch_path), "missing ResNet R2 launch lock"
    launch = read_json(launch_path)
    assert launch["stage"] == "r2"
    assert launch["selector"] == selector
    assert launch["registry_sha256"] == registry_sha256(jobs)
    assert launch["num_jobs"] == 27
    assert launch["allowed_physical_gpus"] == [5, 7]
    source = launch["source"]
    records = [audit_job(job, source, selector) for job in jobs]
    assert len(records) == 27
    assert len({record["cell_id"] for record in records}) == 27
    assert {
        (record["method"], record["depth"]) for record in records
    } == {(method, depth) for method in METHODS for depth in DEPTHS}
    for depth in DEPTHS:
        counts = {
            record["forward_parameter_count"]
            for record in records
            if record["depth"] == depth
            and record["forward_parameter_count"] is not None
        }
        assert len(counts) <= 1, f"forward parameter mismatch at depth {depth}"
    failures = [
        record["cell_id"] for record in records if not record["finite"]
    ]
    report = {
        "audit_status": "passed",
        "stage": "resnet_crossover_r2",
        "complete_grid": True,
        "failure_retaining": True,
        "num_expected_cells": 27,
        "num_audited_cells": len(records),
        "num_finite_cells": len(records) - len(failures),
        "failed_or_incomplete_cells": failures,
        "test_policy": "none",
        "source": source,
        "selector": selector,
        "launch_lock": {
            "path": os.path.relpath(launch_path, ROOT),
            "sha256": sha256(launch_path),
            "registry_sha256": launch["registry_sha256"],
        },
        "records": records,
    }
    os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True)
    with open(args.out, "w", encoding="utf-8") as handle:
        json.dump(report, handle, indent=2, sort_keys=True)
        handle.write("\n")
    print(
        json.dumps(
            {
                "audit_status": report["audit_status"],
                "num_audited_cells": len(records),
                "num_finite_cells": report["num_finite_cells"],
                "failed_or_incomplete_cells": failures,
            },
            indent=2,
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()