"""Step — compute config: cpu_threads resolution (auto / explicit cap / clamp).""" from __future__ import annotations import os from common.runtime import _available_cores, configure_compute, onnx_session_options, resolve_cpu_threads def test_auto_resolves_to_available_cores() -> None: assert resolve_cpu_threads(0) == _available_cores() assert resolve_cpu_threads(-3) == _available_cores() # negative treated as auto, same as 0 def test_explicit_cap_below_available_is_kept() -> None: cap = max(1, _available_cores() - 1) if _available_cores() > 1 else 1 assert resolve_cpu_threads(cap) == cap def test_cap_above_available_is_clamped_down() -> None: huge = _available_cores() + 1000 assert resolve_cpu_threads(huge) == _available_cores() def test_configure_compute_sets_env_vars_to_resolved_value(monkeypatch) -> None: for var in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS", "NUMEXPR_NUM_THREADS", "POLARS_MAX_THREADS", "TOKENIZERS_PARALLELISM"): monkeypatch.delenv(var, raising=False) configure_compute(2) assert os.environ["OMP_NUM_THREADS"] == str(resolve_cpu_threads(2)) assert os.environ["POLARS_MAX_THREADS"] == str(resolve_cpu_threads(2)) assert os.environ["TOKENIZERS_PARALLELISM"] == "true" def test_configure_compute_is_idempotent_setdefault(monkeypatch) -> None: monkeypatch.setenv("OMP_NUM_THREADS", "7") # pre-existing value must survive (setdefault, not overwrite) configure_compute(2) assert os.environ["OMP_NUM_THREADS"] == "7" def test_onnx_session_options_zero_is_auto() -> None: options = onnx_session_options(0) assert options.intra_op_num_threads == 0 def test_onnx_session_options_cap_is_passed_through() -> None: options = onnx_session_options(3) assert options.intra_op_num_threads == 3