bundle/v0/scripts/analyze_full_matrix.py

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path

from loss_aware_dro_repro.core import ContractError, canonical_bytes
from loss_aware_dro_repro.full_matrix_analysis import analyze_full_matrix_run
from loss_aware_dro_repro.reconciled_analysis import analyze_reconciled_full_matrix_run


def _write_immutable(path: Path, value: dict) -> None:
    payload = canonical_bytes(value) + b"\n"
    path.parent.mkdir(parents=True, exist_ok=True)
    if path.exists():
        if path.read_bytes() != payload:
            raise ContractError(f"analysis output already exists with different bytes: {path}")
        return
    descriptor = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL)
    with os.fdopen(descriptor, "wb") as handle:
        handle.write(payload)
        handle.flush()
        os.fsync(handle.fileno())


def main() -> int:
    parser = argparse.ArgumentParser(
        description=(
            "Aggregate ordinary validated 14,000-task outputs and analyze the frozen "
            "capped-5,000 matrix without treating cap hits as convergence."
        )
    )
    parser.add_argument("--manifest", type=Path, required=True)
    source = parser.add_mutually_exclusive_group(required=True)
    source.add_argument(
        "--output-root",
        type=Path,
        help="ordinary batch output root; validates and runs the ordinary aggregator",
    )
    source.add_argument(
        "--reconciled-root",
        type=Path,
        help="sealed recovery bundle root; validates its v2 aggregate without reaggregation",
    )
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--gradient-receipt", type=Path)
    parser.add_argument("--theorem-receipt", type=Path)
    args = parser.parse_args()
    if args.reconciled_root is not None:
        analysis = analyze_reconciled_full_matrix_run(
            args.manifest,
            args.reconciled_root,
            gradient_receipt_path=args.gradient_receipt,
            theorem_receipt_path=args.theorem_receipt,
        )
    else:
        analysis = analyze_full_matrix_run(
            args.manifest,
            args.output_root,
            gradient_receipt_path=args.gradient_receipt,
            theorem_receipt_path=args.theorem_receipt,
        )
    _write_immutable(args.output, analysis)
    print(
        json.dumps(
            {
                "analysis_payload_sha256": analysis["analysis_payload_sha256"],
                "complete": analysis["completeness"]["complete"],
                "validated_success_count": analysis["bindings"][
                    "validated_success_count"
                ],
                "censored_at_5000": analysis["stopping"]["overall"]["counts"][
                    "censored_at_5000"
                ],
                "claim_verdicts": analysis["claim_verdicts"],
                "output": str(args.output.resolve()),
            },
            indent=2,
            sort_keys=True,
        )
    )
    return 0


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