bundle/v0/scripts/run_local_batch.py

from __future__ import annotations

import argparse
import json
from pathlib import Path

from loss_aware_dro_repro.batch_control import (
    aggregate_validated_successes,
    build_batch_manifest,
    build_stopping_sample_batch_manifest,
    run_local_batch,
    write_batch_plan,
)
from loss_aware_dro_repro.core import load_json
from loss_aware_dro_repro.precision_followup import (
    build_precision_followup_batch_manifest,
)


def _result_exit_code(command: str, result: dict) -> int:
    if command == "run":
        blocking_statuses = {
            "failed",
            "failed_validation",
            "skipped_failed_identity",
            "skipped_reserved_identity",
        }
        status_counts = result.get("status_counts", {})
        incomplete = (
            result.get("task_count") != result.get("manifest_task_count")
            or result.get("unissued_task_count") != 0
        )
        if (
            result.get("partial_blocker") is True
            or incomplete
            or any(status_counts.get(status, 0) for status in blocking_statuses)
        ):
            return 2
    if command == "aggregate" and result.get("complete") is not True:
        return 2
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(description="Plan, run, or aggregate an exact-once local batch.")
    commands = parser.add_subparsers(dest="command", required=True)
    plan = commands.add_parser("plan")
    plan.add_argument("--manifest", type=Path, required=True)
    plan.add_argument("--receipt", type=Path, required=True)
    plan.add_argument("--suite", action="append")
    plan.add_argument("--task-id", action="append")
    plan.add_argument("--shard-count", type=int, default=1)
    plan.add_argument("--shard-index", type=int, default=0)
    plan.add_argument("--workers", type=int, default=1)
    plan.add_argument("--execution-config", type=Path)
    stopping_plan = commands.add_parser(
        "plan-stopping-sample",
        help="Plan the frozen 400-task stopping/storage sample with 8 workers.",
    )
    stopping_plan.add_argument("--manifest", type=Path, required=True)
    stopping_plan.add_argument("--receipt", type=Path, required=True)
    followup_plan = commands.add_parser(
        "plan-precision-followup",
        help="Plan the frozen independent 2,580-task precision follow-up.",
    )
    followup_plan.add_argument("--manifest", type=Path, required=True)
    followup_plan.add_argument("--receipt", type=Path, required=True)
    run = commands.add_parser("run")
    run.add_argument("--manifest", type=Path, required=True)
    run.add_argument("--output-root", type=Path, required=True)
    aggregate = commands.add_parser("aggregate")
    aggregate.add_argument("--manifest", type=Path, required=True)
    aggregate.add_argument("--output-root", type=Path, required=True)
    args = parser.parse_args()
    if args.command == "plan":
        requested = load_json(args.execution_config) if args.execution_config else None
        manifest, receipt = build_batch_manifest(
            suites=args.suite,
            task_ids=args.task_id,
            shard_count=args.shard_count,
            shard_index=args.shard_index,
            worker_count=args.workers,
            requested_execution_config=requested,
        )
        write_batch_plan(args.manifest, args.receipt, manifest, receipt)
        result = receipt
    elif args.command == "plan-stopping-sample":
        manifest, receipt = build_stopping_sample_batch_manifest()
        write_batch_plan(args.manifest, args.receipt, manifest, receipt)
        result = receipt
    elif args.command == "plan-precision-followup":
        manifest, receipt = build_precision_followup_batch_manifest()
        write_batch_plan(args.manifest, args.receipt, manifest, receipt)
        result = receipt
    elif args.command == "run":
        result = run_local_batch(args.manifest, args.output_root)
    else:
        result = aggregate_validated_successes(args.manifest, args.output_root)
    print(json.dumps(result, indent=2, sort_keys=True))
    return _result_exit_code(args.command, result)


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