| """Threading determinism: pin BLAS / OpenMP / pyarrow to single-threaded mode. | |
| Multi-threaded floating-point reductions reorder operands non-deterministically | |
| on each call, breaking the byte-identity guarantee in AGENTS.md §4 rule 3. Each | |
| pipeline calls `pin_threads()` at import time to lock the process to a single | |
| thread before any numerical work runs. | |
| Honors pre-set env vars: if the caller exported `OMP_NUM_THREADS=4` upstream, | |
| that value is preserved (we use `setdefault`, not `setitem`). The user is | |
| responsible for the determinism trade-off in that case. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import pyarrow as pa | |
| _ENV_VARS: tuple[str, ...] = ( | |
| "OMP_NUM_THREADS", | |
| "OPENBLAS_NUM_THREADS", | |
| "MKL_NUM_THREADS", | |
| ) | |
| def pin_threads() -> None: | |
| """Pin BLAS / OpenMP / pyarrow to single-threaded mode (idempotent).""" | |
| for var in _ENV_VARS: | |
| os.environ.setdefault(var, "1") | |
| pa.set_cpu_count(1) | |
| pa.set_io_thread_count(1) | |