bundle/v0/scripts/recover_failed_batch.py

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

LANE_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(LANE_ROOT / "src"))

from loss_aware_dro_repro.batch_recovery import build_recovery_manifest, reconcile_aggregate, run_recovery


def main() -> int:
    parser = argparse.ArgumentParser(description="Recover explicitly recorded controller exceptions without changing the scientific task contract.")
    commands = parser.add_subparsers(dest="command", required=True)
    plan = commands.add_parser("plan")
    plan.add_argument("--source-manifest", type=Path, required=True)
    plan.add_argument("--source-output-root", type=Path, required=True)
    plan.add_argument("--source-lane", type=Path, required=True)
    plan.add_argument("--recovery-lane", type=Path, required=True)
    plan.add_argument("--python", type=Path, required=True)
    plan.add_argument("--recovery-output-root", type=Path, required=True)
    plan.add_argument("--recovery-manifest", type=Path, required=True)
    run = commands.add_parser("run")
    run.add_argument("--recovery-manifest", type=Path, required=True)
    run.add_argument("--recovery-output-root", type=Path, required=True)
    aggregate = commands.add_parser("aggregate")
    aggregate.add_argument("--recovery-manifest", type=Path, required=True)
    aggregate.add_argument("--recovery-output-root", type=Path, required=True)
    aggregate.add_argument("--aggregate-root", type=Path, required=True)
    args = parser.parse_args()
    if args.command == "plan":
        result = build_recovery_manifest(
            args.source_manifest,
            args.source_output_root,
            args.source_lane,
            args.recovery_lane,
            python_executable=args.python,
            recovery_output_root=args.recovery_output_root,
        )
        args.recovery_manifest.parent.mkdir(parents=True, exist_ok=True)
        if args.recovery_manifest.exists():
            raise FileExistsError(args.recovery_manifest)
        args.recovery_manifest.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8", newline="\n")
    elif args.command == "run":
        result = run_recovery(args.recovery_manifest, args.recovery_output_root)
    else:
        result = reconcile_aggregate(args.recovery_manifest, args.recovery_output_root, args.aggregate_root)
    print(json.dumps(result, indent=2, sort_keys=True))
    return 0


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