bundle/v0/scripts/local_batch_worker.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.core import load_json
from loss_aware_dro_repro.task_dispatch import run_task


def main() -> int:
    parser = argparse.ArgumentParser(description="One-thread exact-once local scientific task worker.")
    parser.add_argument("--task-id", required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--execution-config", type=Path)
    args = parser.parse_args()
    requested = load_json(args.execution_config) if args.execution_config else None
    result = run_task(args.task_id, args.output_dir, requested_execution_config=requested)
    print(json.dumps({"task_id": result["task_id"], "run_identity": result["run_identity"], "status": result["status"]}, sort_keys=True))
    return 0


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