bundle/v0/scripts/run_checkpoint_resume_canary.py

from __future__ import annotations

import os


for _name in (
    "OMP_NUM_THREADS",
    "OPENBLAS_NUM_THREADS",
    "MKL_NUM_THREADS",
    "VECLIB_MAXIMUM_THREADS",
    "NUMEXPR_NUM_THREADS",
    "BLIS_NUM_THREADS",
):
    os.environ[_name] = "1"


import argparse
import json
from pathlib import Path

from loss_aware_dro_repro.checkpoint_resume import (
    RESUMABLE_EXIT_CODE,
    refuse_same_identity_restart,
    resume_canary,
    run_canary,
    run_resume_worker,
    run_start_worker,
    start_canary,
    validate_canary,
)


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Run the local, claim-ineligible Algorithm 2 checkpoint/resume canary."
    )
    commands = parser.add_subparsers(dest="command", required=True)
    for name in ("run", "start"):
        command = commands.add_parser(name)
        command.add_argument(
            "--config",
            type=Path,
            default=Path("configs/checkpoint_resume_canary_v1.json"),
        )
        command.add_argument("--output-root", type=Path, required=True)
    for name in ("resume", "validate"):
        command = commands.add_parser(name)
        command.add_argument("--output-root", type=Path, required=True)
    for name in ("worker-start", "worker-resume", "restart"):
        command = commands.add_parser(name)
        command.add_argument("--output-root", type=Path, required=True)
        command.add_argument("--task-index", type=int, required=True)
        if name != "restart":
            command.add_argument("--launch-binding-hash", required=True)
    args = parser.parse_args()
    if args.command == "run":
        result = run_canary(args.config, args.output_root)
    elif args.command == "start":
        result = start_canary(args.config, args.output_root)
    elif args.command == "resume":
        result = resume_canary(args.output_root)
    elif args.command == "validate":
        result = validate_canary(args.output_root)
    elif args.command == "worker-start":
        result = run_start_worker(
            args.output_root, args.task_index, args.launch_binding_hash
        )
        print(json.dumps(result, sort_keys=True))
        return RESUMABLE_EXIT_CODE
    elif args.command == "worker-resume":
        result = run_resume_worker(
            args.output_root, args.task_index, args.launch_binding_hash
        )
    else:
        refuse_same_identity_restart(args.output_root, args.task_index)
        raise AssertionError("same-identity restart unexpectedly returned")
    print(json.dumps(result, indent=2, sort_keys=True))
    return 0


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