bundle/v0/scripts/run_optimizer_probe.py

from __future__ import annotations

import argparse
import os
from pathlib import Path


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"
os.environ["LOSS_AWARE_OPTIMIZER_PROBE_CLI_BOOTSTRAP"] = (
    "before_scientific_imports_v1"
)

from loss_aware_dro_repro.optimizer_probe import run_optimizer_probe  # noqa: E402


def main() -> None:
    parser = argparse.ArgumentParser(
        description=(
            "Run the local-only, claim-ineligible optimizer-fidelity probe exactly once."
        )
    )
    parser.add_argument("--config", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    arguments = parser.parse_args()
    run_optimizer_probe(arguments.config, arguments.output)


if __name__ == "__main__":
    main()