stream-limits 0.2.1__tar.gz → 0.3.0__tar.gz

{stream_limits-0.2.1 → stream_limits-0.3.0}/CHANGELOG.md

@@ -1,5 +1,28 @@
 # Changelog
 
+## 0.3.0
+
+- **A new function, `set_openmp_wait_policy`, controls how idle OpenMP threads
+  wait.** Programs that make many short parallel calls, such as numba loops run
+  once per batch of data, can run noticeably slower when the idle threads spin
+  while they wait for the next call, especially on processors that mix fast and
+  efficient cores. Telling the threads to sleep instead (`OMP_WAIT_POLICY=passive`,
+  the default here) fixed this in diff-stream: fits ran 2% to 15% faster on a
+  cloud server, and on a laptop with mixed cores 8 threads went from slower than
+  one thread to faster. Programs with a few long parallel calls, such as
+  hdfe-stream and logit-stream, run the same either way. On Windows, where
+  numba's OpenMP runtime ignores the setting, the call changes nothing.
+- **A package has to ask for it.** Importing stream-limits, and
+  `apply_scheduler_defaults`, still change nothing, because the setting reaches
+  every OpenMP library in the process (BLAS builds on OpenMP, scikit-learn,
+  LightGBM). A value you have set yourself is kept unless you pass
+  `override=True`.
+- **It must run before OpenMP starts.** The setting is read once, when the
+  OpenMP runtime loads; numba loads its runtime at its first parallel call, so
+  calling the function after importing numba is still in time. The function
+  reports when a runtime was already loaded (Linux), in which case the setting
+  reaches only programs started afterwards.
+
 ## 0.2.1
 
 - **The watchdog no longer counts the pages of memory-mapped files.** A program

{stream_limits-0.2.1 → stream_limits-0.3.0}/CITATION.cff

@@ -7,8 +7,8 @@ authors:
   - family-names: Tucker
     given-names: Lee
     email: lee.c.tucker@census.gov
-version: 0.2.1
-date-released: 2026-11-04
+version: 0.3.0
+date-released: 2026-11-08
 license: MIT
 repository-code: "https://github.com/leetucker/stream-limits"
 url: "https://github.com/leetucker/stream-limits"

{stream_limits-0.2.1/stream_limits.egg-info → stream_limits-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stream-limits
-Version: 0.2.1
+Version: 0.3.0
 Summary: Cap the threads and memory of numeric Python libraries (Polars, numba, BLAS, pyarrow) at a scheduler's allocation, and measure peak memory
 Author: Lee Tucker
 License-Expression: MIT

{stream_limits-0.2.1 → stream_limits-0.3.0}/docs/reference.md

@@ -15,6 +15,8 @@
 
 .. autofunction:: thread_scope
 
+.. autofunction:: set_openmp_wait_policy
+
 .. autofunction:: polars_spills
 ```
 

{stream_limits-0.2.1 → stream_limits-0.3.0}/docs/threads.md

@@ -81,3 +81,23 @@ allocation, read from the environment, cgroups or the scheduler (see
 [Schedulers and containers](schedulers.md)), and otherwise the machine's. A job that runs several
 processes should give each its share, for example
 `STREAM_LIMITS_MAX_THREADS=$(( NCPUS / 4 ))` for four.
+
+
+## How OpenMP threads wait
+
+OpenMP threads that finish a parallel region either spin, waiting for the next one, or
+sleep. Spinning is quicker to resume, but a program that makes many short parallel calls
+(numba's `prange` loops, one per batch of data) can run several times slower with it,
+notably on CPUs that mix performance and efficiency cores, and a little slower on others.
+Programs with few, long parallel regions are unaffected: hdfe-stream's solvers and
+logit-stream's fits ran as fast either way.
+
+{py:func}`stream_limits.set_openmp_wait_policy` sets `OMP_WAIT_POLICY` (to `"passive"` unless you
+say `"active"`). A package opts in by calling it, because the setting reaches every
+OpenMP library in the process (BLAS builds on OpenMP, scikit-learn, LightGBM); importing
+stream-limits, and `apply_scheduler_defaults`, leave it alone. A value you have set is kept unless
+you pass `override=True`. The runtime reads the variable when it loads, which for numba is
+at its first parallel call, so call the function before that. The result says when a runtime
+was already loaded (Linux), in which case the variable is set for child processes only.
+On Windows, numba uses Microsoft's OpenMP runtime, which does not read the variable, so the
+call changes nothing there.

{stream_limits-0.2.1 → stream_limits-0.3.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "stream-limits"
-version = "0.2.1"
+version = "0.3.0"
 description = "Cap the threads and memory of numeric Python libraries (Polars, numba, BLAS, pyarrow) at a scheduler's allocation, and measure peak memory"
 readme = "README.md"
 license = "MIT"

{stream_limits-0.2.1 → stream_limits-0.3.0}/stream_limits/__init__.py

@@ -87,10 +87,11 @@ import warnings
 from pathlib import Path
 from typing import Iterator, Mapping
 
-__version__ = "0.2.1"
+__version__ = "0.3.0"
 
 __all__ = [
     "set_limits", "limit_threads", "limit_memory", "apply_scheduler_defaults", "thread_scope",
+    "set_openmp_wait_policy",
     "detect_cpus", "detect_memory_gb", "describe", "MemoryBudget",
     "peak_rss_gb", "reset_peak_rss", "current_rss_gb",
 ]
@@ -790,6 +791,51 @@ def thread_scope(n: int | None = None) -> Iterator[None]:
         numba.set_num_threads(previous)
 
 
+def _openmp_loaded() -> bool | None:
+    """Whether an OpenMP runtime is already in this process (Linux only; None
+    elsewhere). It reads ``OMP_WAIT_POLICY`` once, when it loads."""
+    try:
+        maps = Path("/proc/self/maps").read_text()
+    except OSError:
+        return None
+    return any(name in maps for name in ("libgomp", "libiomp5", "libomp"))
+
+
+def set_openmp_wait_policy(policy: str = "passive", *, override: bool = False,
+                           environ=None) -> dict:
+    """Set how OpenMP threads wait between parallel regions, by ``OMP_WAIT_POLICY``.
+
+    With ``"passive"`` an idle thread sleeps until the next region; with
+    ``"active"`` it spins. A package whose compiled code makes many short
+    parallel calls (numba's ``prange`` loops, one per batch of data) can run
+    several times slower with spinning threads, notably on CPUs with
+    performance and efficiency cores, and a little slower on others. Code with
+    few, long parallel regions (hdfe-stream's solvers, logit-stream's fits)
+    is unaffected, and this function is not called by importing the package:
+    a package opts in, since the setting reaches every OpenMP library in the
+    process (BLAS builds on OpenMP, scikit-learn, LightGBM). Microsoft's OpenMP
+    runtime, the one numba uses on Windows, does not read the variable, so there
+    the call changes nothing.
+
+    A value the user has set is left alone unless `override`. The runtime reads
+    the variable once, when it loads (numba loads its OpenMP layer at the first
+    parallel call), so call this before the first parallel call; the result's
+    ``"late"`` is True when a runtime was already loaded (Linux), in which case
+    the variable is still set, for child processes, but this process is
+    unaffected. Returns ``{"policy": ..., "set": bool, "late": bool | None}``,
+    where ``"policy"`` is the value in effect."""
+    if policy not in ("passive", "active"):
+        raise ValueError(f"policy must be 'passive' or 'active', not {policy!r}")
+    environ = os.environ if environ is None else environ
+    current = environ.get("OMP_WAIT_POLICY")
+    if current and not override:
+        return {"policy": current.lower(), "set": False, "late": None}
+    late = _openmp_loaded() if environ is os.environ else None
+    environ["OMP_WAIT_POLICY"] = policy
+    return {"policy": policy, "set": True, "late": late}
+
+
+
 # ----------------------------------------------------- peak memory used
 def _windows_memory(handle=None) -> tuple[int, int] | None:
     """This process's ``(working set, peak working set)`` in bytes, from the

{stream_limits-0.2.1 → stream_limits-0.3.0/stream_limits.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stream-limits
-Version: 0.2.1
+Version: 0.3.0
 Summary: Cap the threads and memory of numeric Python libraries (Polars, numba, BLAS, pyarrow) at a scheduler's allocation, and measure peak memory
 Author: Lee Tucker
 License-Expression: MIT

{stream_limits-0.2.1 → stream_limits-0.3.0}/tests/test_override_and_libraries.py

@@ -213,3 +213,68 @@ def test_a_preset_spill_dir_is_kept_unless_override(tmp_path):
     a = _run(code.format(extra=""), POLARS_OOC_SPILL_DIR="/somewhere/else")
     b = _run(code.format(extra=", override=True"), POLARS_OOC_SPILL_DIR="/somewhere/else")
     assert a.stdout.strip() == "/somewhere/else" and b.stdout.strip() == str(tmp_path), (a.stderr, b.stderr)
+
+
+# --- set_openmp_wait_policy ---
+
+def test_wait_policy_sets_the_variable_and_respects_the_user():
+    env = {}
+    assert sl.set_openmp_wait_policy(environ=env) == {"policy": "passive", "set": True, "late": None}
+    assert env["OMP_WAIT_POLICY"] == "passive"
+    env = {"OMP_WAIT_POLICY": "ACTIVE"}
+    assert sl.set_openmp_wait_policy(environ=env) == {"policy": "active", "set": False, "late": None}
+    assert env["OMP_WAIT_POLICY"] == "ACTIVE"
+    assert sl.set_openmp_wait_policy(environ=env, override=True)["set"] is True
+    assert env["OMP_WAIT_POLICY"] == "passive"
+    with pytest.raises(ValueError):
+        sl.set_openmp_wait_policy("sleepy", environ={})
+
+
+def test_importing_the_module_leaves_the_policy_alone():
+    r = _run("import os, stream_limits; print(os.environ.get('OMP_WAIT_POLICY'))")
+    assert r.stdout.strip() == "None"
+    r = _run("import os, stream_limits; stream_limits.apply_scheduler_defaults(); "
+             "print(os.environ.get('OMP_WAIT_POLICY'))", NCPUS="2")
+    assert r.stdout.strip() == "None"
+
+
+NUMBA_PROBE = """
+import numpy as np, sys
+{before}
+import numba
+from numba import njit, prange
+import stream_limits
+print(stream_limits.set_openmp_wait_policy()["late"])
+@njit(parallel=True)
+def f(a):
+    s = 0.0
+    for i in prange(a.size):
+        s += a[i]
+    return s
+f(np.ones(100))
+{after}
+"""
+
+
+def test_wait_policy_takes_effect_when_numba_is_imported_first():
+    # numba loads its OpenMP runtime at the first parallel call, which is after
+    r = _run(NUMBA_PROBE.format(before="", after="print(numba.threading_layer())"),
+             OMP_DISPLAY_ENV="verbose")
+    if "omp" not in r.stdout.split():
+        pytest.skip("numba's OpenMP threading layer is not in use")
+    # "late" is only detectable on Linux (it reads /proc/self/maps)
+    assert r.stdout.split()[0] == ("False" if sys.platform.startswith("linux") else "None")
+    if sys.platform == "win32":
+        pytest.skip("Microsoft's OpenMP runtime does not read or display OMP_WAIT_POLICY")
+    assert "OMP_WAIT_POLICY = 'PASSIVE'" in r.stderr
+
+
+def test_wait_policy_reports_a_runtime_that_is_already_loaded():
+    r = _run(NUMBA_PROBE.format(before="", after="") +
+             "\nimport stream_limits as s\nprint(s.set_openmp_wait_policy(override=True)['late'])")
+    if not sys.platform.startswith("linux"):
+        pytest.skip("detection needs /proc/self/maps")
+    lines = r.stdout.split()
+    if lines[-1] == "False":
+        pytest.skip("numba's OpenMP layer is not in use")
+    assert lines[-1] == "True"