bundle/v0/scripts/run_oos_timing_sample.py

from __future__ import annotations

import argparse
import os
from pathlib import Path


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

from loss_aware_dro_repro.oos_timing_harness import run_oos_timing_sample  # noqa: E402


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Time the local grouped 10M absolute and squared regression OOS contracts."
    )
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    receipt = run_oos_timing_sample(args.output)
    print(receipt["run_identity"])
    print(args.output.resolve())


if __name__ == "__main__":
    main()