bundle/v0/scripts/run_cpu_timing_sample.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_TIMING_CLI_BOOTSTRAP"] = "before_scientific_imports_v1"

from loss_aware_dro_repro.timing_harness import run_timing_sample  # noqa: E402


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Run the local-only, claim-ineligible CPU timing sample."
    )
    parser.add_argument("--config", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    arguments = parser.parse_args()
    run_timing_sample(arguments.config, arguments.output)


if __name__ == "__main__":
    main()