scripts/launch_job.py

#!/usr/bin/env python3
"""Prepare, submit, reconcile, poll, or cancel one durable direct HF Job attempt."""

from __future__ import annotations

import argparse
import json
import sys
from datetime import UTC, datetime
from pathlib import Path

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

from repro_control.hf_provider import HFJobsProvider
from repro_control.launcher import (
    AmbiguousSubmission,
    BudgetSnapshot,
    DirectLauncher,
    LaunchError,
    LauncherState,
)


def _launcher(args) -> DirectLauncher:
    provider = HFJobsProvider(namespace=args.namespace) if args.action != "prepare" else None
    return DirectLauncher(
        args.state,
        provider=provider,
        enable_submit=args.enable_submit,
    )


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--state", type=Path, required=True)
    parser.add_argument("--namespace", default="Mindcraft")
    parser.add_argument(
        "--enable-submit",
        action="store_true",
        help="Required for submit/cancel; omitted by default to keep the CLI non-mutating.",
    )
    sub = parser.add_subparsers(dest="action", required=True)
    prepare = sub.add_parser("prepare")
    prepare.add_argument("--request", type=Path, required=True)
    prepare.add_argument("--budget-snapshot", type=Path, required=True)
    prepare.add_argument("--guarded-critical-path-seconds", type=int, required=True)
    submit = sub.add_parser("submit")
    submit.add_argument("--attempt-id", required=True)
    reconcile = sub.add_parser("reconcile")
    reconcile.add_argument("--attempt-id", required=True)
    poll = sub.add_parser("poll")
    poll.add_argument("--attempt-id", required=True)
    cancel = sub.add_parser("cancel")
    cancel.add_argument("--attempt-id", required=True)
    sub.add_parser("show")
    args = parser.parse_args()
    launcher = _launcher(args)
    try:
        if args.action == "prepare":
            request = json.loads(args.request.read_text())
            budget = BudgetSnapshot(**json.loads(args.budget_snapshot.read_text()))
            row = launcher.prepare(
                request,
                budget=budget,
                now=datetime.now(UTC),
                guarded_critical_path_seconds=args.guarded_critical_path_seconds,
            )
        elif args.action == "submit":
            if not args.enable_submit:
                raise LaunchError("submit requires --enable-submit")
            row = launcher.submit_prepared(args.attempt_id)
        elif args.action == "reconcile":
            row = launcher.reconcile_ambiguous(args.attempt_id)
        elif args.action == "poll":
            row = launcher.poll(args.attempt_id, now=datetime.now(UTC))
        elif args.action == "cancel":
            if not args.enable_submit:
                raise LaunchError("cancel requires --enable-submit")
            launcher.cancel(args.attempt_id)
            row = next(
                item
                for item in LauncherState.load(args.state).attempts
                if item["attempt_id"] == args.attempt_id
            )
        else:
            print(
                json.dumps(
                    {
                        "format": 1,
                        "attempts": LauncherState.load(args.state).attempts,
                    },
                    sort_keys=True,
                )
            )
            return 0
    except AmbiguousSubmission as exc:
        print(str(exc), file=sys.stderr)
        return 3
    except (LaunchError, ValueError, KeyError) as exc:
        print(str(exc), file=sys.stderr)
        return 2
    print(json.dumps(row, sort_keys=True))
    return 0


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