path: root/experiments/import_native_baseline.py

#!/usr/bin/env python3
"""Import completed author-code runs into deterministic, audited JSON records.

Run this script with the baseline's own Python environment. BurstCCN needs the
``wandb`` package to decode its offline event file; Dual Prop needs ``numpy``
to read the author's ``hist.npy``. The importer never evaluates or retrains a
model and refuses incomplete histories, dirty source trees, revision drift,
non-finite metrics, and protocol/config mismatches.
"""
import argparse
import hashlib
import json
import math
import os
import subprocess
import sys


ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEFAULT_PROTOCOLS = os.path.join(
    ROOT, "experiments", "native_baseline_protocols.json")


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


def run_text(command, cwd=None):
    return subprocess.run(
        command, cwd=cwd, check=True, capture_output=True, text=True).stdout.strip()


def git_record(path):
    return {
        "path": os.path.abspath(path),
        "revision": run_text(["git", "rev-parse", "HEAD"], cwd=path),
        "tracked_dirty": bool(run_text(
            ["git", "status", "--porcelain", "--untracked-files=no"], cwd=path)),
    }


def finite_number(value, label):
    value = float(value)
    if not math.isfinite(value):
        raise ValueError(f"{label} is non-finite: {value!r}")
    return value


def finite_list(values, label):
    return [finite_number(value, f"{label}[{index}]")
            for index, value in enumerate(values)]


def summarize_burst_history(rows, expected_epochs):
    by_epoch = {}
    for row in rows:
        if "epoch" not in row:
            continue
        epoch = int(row["epoch"])
        if "epoch/top1_error/test" in row:
            by_epoch[epoch] = row
    expected = list(range(expected_epochs))
    observed = sorted(by_epoch)
    if observed != expected:
        raise ValueError(
            f"BurstCCN history incomplete: expected epochs 0..{expected_epochs - 1}, "
            f"found {len(observed)} rows ending at {observed[-1] if observed else None}")

    test_errors = finite_list(
        [by_epoch[e]["epoch/top1_error/test"] for e in expected], "test top1 error")
    val_errors = finite_list(
        [by_epoch[e]["epoch/top1_error/val"] for e in expected], "val top1 error")
    final = expected[-1]
    best_test = min(range(expected_epochs), key=lambda e: test_errors[e])
    best_val = min(range(expected_epochs), key=lambda e: val_errors[e])
    runtime = finite_number(by_epoch[final].get("_runtime"), "W&B runtime")
    return {
        "epochs_observed": expected_epochs,
        "optimization_epochs": expected_epochs - 1,
        "initial_evaluation_epoch": 0,
        "primary_final_test_accuracy_percent": 100.0 - test_errors[final],
        "primary_final_epoch": final,
        "final_validation_accuracy_percent": 100.0 - val_errors[final],
        "validation_selected_epoch": best_val,
        "validation_selected_validation_accuracy_percent": 100.0 - val_errors[best_val],
        "validation_selected_test_accuracy_percent": 100.0 - test_errors[best_val],
        "test_selected_epoch": best_test,
        "test_selected_test_accuracy_percent": 100.0 - test_errors[best_test],
        "test_evaluations": expected_epochs,
        "wall_s": runtime,
        "wall_definition": "W&B runtime from run initialization through the final epoch log",
    }


def read_burst_history(path):
    try:
        from wandb.proto import wandb_internal_pb2
        from wandb.sdk.internal.datastore import DataStore
    except ImportError as exc:
        raise RuntimeError(
            "BurstCCN import must run in an environment containing wandb") from exc

    store = DataStore()
    store.open_for_scan(path)
    rows = []
    while True:
        data = store.scan_data()
        if data is None:
            break
        record = wandb_internal_pb2.Record()
        record.ParseFromString(data)
        if record.WhichOneof("record_type") != "history":
            continue
        row = {}
        for item in record.history.item:
            key = item.key or "/".join(item.nested_key)
            if key:
                row[key] = json.loads(item.value_json)
        rows.append(row)
    return rows


def summarize_dualprop_hist(hist, expected_epochs):
    required_vectors = (
        "train_loss", "train_accuracy", "train_time", "val_loss",
        "val_accuracy", "val_top5accuracy", "val_time")
    vectors = {}
    for key in required_vectors:
        if key not in hist:
            raise ValueError(f"Dual Prop history lacks {key!r}")
        values = list(hist[key])
        if len(values) != expected_epochs:
            raise ValueError(
                f"Dual Prop {key} has {len(values)} epochs, expected {expected_epochs}")
        vectors[key] = finite_list(values, key)
    for key in ("test_loss", "test_accuracy", "test_top5accuracy", "test_time"):
        if key not in hist:
            raise ValueError(f"Dual Prop history lacks {key!r}")
    test_accuracy = finite_number(hist["test_accuracy"], "test_accuracy")
    test_top5 = finite_number(hist["test_top5accuracy"], "test_top5accuracy")
    test_loss = finite_number(hist["test_loss"], "test_loss")
    test_time = finite_number(hist["test_time"], "test_time")
    best_val = max(range(expected_epochs), key=lambda e: vectors["val_accuracy"][e])
    component_wall = (sum(vectors["train_time"]) + sum(vectors["val_time"])
                      + test_time)
    return {
        "epochs_observed": expected_epochs,
        "final_epoch": expected_epochs,
        "final_train_accuracy_percent": vectors["train_accuracy"][-1],
        "final_validation_accuracy_percent": vectors["val_accuracy"][-1],
        "validation_selected_epoch": best_val + 1,
        "validation_selected_validation_accuracy_percent": vectors["val_accuracy"][best_val],
        "validation_selected_test_accuracy_percent": test_accuracy,
        "primary_test_accuracy_percent": test_accuracy,
        "primary_test_top5_accuracy_percent": test_top5,
        "primary_test_loss": test_loss,
        "test_evaluations": 1,
        "test_selected_test_accuracy_percent": None,
        "wall_s": component_wall,
        "training_wall_s": sum(vectors["train_time"]),
        "validation_wall_s": sum(vectors["val_time"]),
        "test_wall_s": test_time,
        "wall_definition": (
            "sum of author-measured train, validation, and final test runtimes; "
            "excludes setup, checkpoint I/O, and diagnostic plotting"),
    }


def read_dualprop_hist(path):
    try:
        import numpy as np
    except ImportError as exc:
        raise RuntimeError(
            "Dual Prop import must run in an environment containing numpy") from exc
    value = np.load(path, allow_pickle=True)
    if value.shape != ():
        raise ValueError(f"expected scalar object hist.npy, found shape {value.shape}")
    hist = value.item()
    if not isinstance(hist, dict):
        raise ValueError(f"expected dict in hist.npy, found {type(hist).__name__}")
    return hist


def nested_value(mapping, dotted_key):
    value = mapping
    for key in dotted_key.split("."):
        if not isinstance(value, dict) or key not in value:
            raise ValueError(f"resolved config lacks {dotted_key!r}")
        value = value[key]
    return value


def read_and_validate_config(path, required):
    try:
        from omegaconf import OmegaConf
    except ImportError as exc:
        raise RuntimeError(
            "resolved Hydra config import must run in the BurstCCN environment") from exc
    config = OmegaConf.to_container(OmegaConf.load(path), resolve=True)
    mismatches = []
    for key, expected in required.items():
        actual = nested_value(config, key)
        if actual != expected:
            mismatches.append(f"{key}: expected {expected!r}, found {actual!r}")
    if mismatches:
        raise ValueError("resolved config mismatch:\n" + "\n".join(mismatches))
    return config


def environment_record(include_freeze=True):
    record = {
        "python_executable": os.path.abspath(sys.executable),
        "python_version": sys.version,
    }
    if include_freeze:
        freeze = run_text([sys.executable, "-m", "pip", "freeze"]).splitlines()
        freeze = sorted(line for line in freeze if line.strip())
        encoded = ("\n".join(freeze) + "\n").encode()
        record["pip_freeze"] = freeze
        record["pip_freeze_sha256"] = hashlib.sha256(encoded).hexdigest()
    return record


def gpu_record(index):
    query = run_text([
        "nvidia-smi", f"--id={index}",
        "--query-gpu=index,name,uuid,driver_version,memory.total",
        "--format=csv,noheader,nounits",
    ])
    fields = [value.strip() for value in query.split(",")]
    if len(fields) != 5:
        raise ValueError(f"unexpected nvidia-smi record: {query!r}")
    return dict(zip(("physical_index", "name", "uuid", "driver_version",
                     "memory_total_mib"), fields))


def load_protocol(path, protocol_id):
    with open(path) as handle:
        ledger = json.load(handle)
    if ledger.get("schema_version") != 1:
        raise ValueError("unsupported native protocol schema")
    try:
        return ledger["protocols"][protocol_id]
    except KeyError as exc:
        raise ValueError(f"unknown protocol {protocol_id!r}") from exc


def importer_provenance():
    return git_record(ROOT)


def import_result(args):
    protocol = load_protocol(args.protocols, args.protocol)
    source = git_record(args.source_dir)
    if source["revision"] != protocol["source_revision"]:
        raise ValueError(
            f"source revision mismatch: expected {protocol['source_revision']}, "
            f"found {source['revision']}")
    if source["tracked_dirty"]:
        raise ValueError("author source tree has tracked modifications")
    local_source = importer_provenance()
    if local_source["tracked_dirty"]:
        raise ValueError("SDIL importer source tree has tracked modifications")

    if protocol["method"] == "BurstCCN":
        if not args.resolved_config:
            raise ValueError("BurstCCN requires --resolved-config")
        resolved = read_and_validate_config(
            args.resolved_config, protocol["required_config"])
        history = read_burst_history(args.artifact)
        metrics = summarize_burst_history(history, protocol["expected_epochs"])
    elif protocol["method"] == "Dual Propagation":
        if args.resolved_config:
            raise ValueError("Dual Prop protocol is frozen directly in the ledger")
        resolved = protocol["resolved_config"]
        history = read_dualprop_hist(args.artifact)
        metrics = summarize_dualprop_hist(history, protocol["expected_epochs"])
    else:
        raise ValueError(f"unsupported native method {protocol['method']!r}")

    artifacts = [file_record(args.artifact)]
    if args.resolved_config:
        artifacts.append(file_record(args.resolved_config))
    dataset_artifacts = [file_record(path) for path in args.dataset_artifact]
    dataset_hashes = {record["sha256"] for record in dataset_artifacts}
    missing_dataset_hashes = (set(protocol["required_dataset_sha256"])
                              - dataset_hashes)
    if missing_dataset_hashes:
        raise ValueError(
            "required dataset artifact hashes are missing: "
            + ", ".join(sorted(missing_dataset_hashes)))
    gpu = gpu_record(args.gpu_index)
    for key, expected in protocol["expected_gpu"].items():
        if gpu.get(key) != expected:
            raise ValueError(
                f"GPU {key} mismatch: expected {expected!r}, found {gpu.get(key)!r}")
    result = {
        "schema_version": 1,
        "protocol_id": args.protocol,
        "method": protocol["method"],
        "dataset": protocol["dataset"],
        "seed": protocol["seed"],
        "protocol": protocol,
        "resolved_config": resolved,
        "metrics": metrics,
        "audit": {
            "complete": True,
            "finite": True,
            "expected_epochs": protocol["expected_epochs"],
            "source_revision_match": True,
            "source_tracked_clean": True,
            "selection_semantics_frozen": True,
        },
        "source_provenance": source,
        "importer_provenance": local_source,
        "runtime_environment": environment_record(not args.skip_package_freeze),
        "gpu": gpu,
        "dataset_artifacts": dataset_artifacts,
        "run_artifacts": artifacts,
    }
    if os.path.exists(args.output) and not args.overwrite:
        raise FileExistsError(f"refusing to overwrite {args.output}")
    os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True)
    with open(args.output, "w") as handle:
        json.dump(result, handle, indent=2, sort_keys=True)
        handle.write("\n")
    print(json.dumps({
        "output": os.path.abspath(args.output),
        "method": result["method"],
        "metrics": result["metrics"],
    }, indent=2, sort_keys=True))


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--protocol", required=True)
    parser.add_argument("--protocols", default=DEFAULT_PROTOCOLS)
    parser.add_argument("--source-dir", required=True)
    parser.add_argument("--artifact", required=True,
                        help="BurstCCN run-*.wandb or Dual Prop hist.npy")
    parser.add_argument("--resolved-config",
                        help="BurstCCN's .hydra/config.yaml")
    parser.add_argument("--dataset-artifact", action="append", default=[])
    parser.add_argument("--gpu-index", required=True, type=int)
    parser.add_argument("--output", required=True)
    parser.add_argument("--skip-package-freeze", action="store_true",
                        help="development only; final audited imports retain the freeze")
    parser.add_argument("--overwrite", action="store_true")
    return parser.parse_args()


if __name__ == "__main__":
    import_result(parse_args())