File size: 1,848 Bytes
a41b023 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | """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
|