bundle/v0/scripts/analyze_precision_followup.py

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path

from loss_aware_dro_repro.core import canonical_bytes, sha256_value
from loss_aware_dro_repro.precision_followup import analyze_precision_followup


def _publish(path: Path, payload: bytes) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(f"{path.name}.{os.getpid()}.tmp")
    temporary.write_bytes(payload)
    os.replace(temporary, path)


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Analyze the frozen independent stopping and storage follow-up."
    )
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    parser.add_argument("--aggregate-summary", type=Path, required=True)
    parser.add_argument("--aggregate-results", type=Path, required=True)
    args = parser.parse_args()
    analysis = analyze_precision_followup(
        batch_manifest_path=args.manifest,
        batch_output_root=args.output_root,
        aggregate_summary_path=args.aggregate_summary,
        aggregate_results_path=args.aggregate_results,
    )
    identity = sha256_value(analysis)
    output = (
        args.output_root
        / "analyses"
        / identity.removeprefix("sha256:")
        / "analysis.json"
    )
    _publish(output, canonical_bytes(analysis))
    print(
        json.dumps(
            {
                "analysis_identity": identity,
                "analysis_path": str(output.resolve()),
                "sample_task_count": analysis["sample_task_count"],
                "precision_pass": analysis["precision_pass"],
                "precision_failures": analysis["precision_failures"],
                "weighted_estimates": analysis["weighted_estimates"],
                "eligibility": analysis["eligibility"],
                "scientific_verdicts": analysis["scientific_verdicts"],
            },
            indent=2,
            sort_keys=True,
        )
    )
    return 0 if analysis["precision_pass"] else 2


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