From 5a39588b4acc8a6beee8c9c53878bcecc620a036 Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Thu, 24 Sep 2026 13:15:37 +0200 Subject: [PATCH 1/3] fix: use the probed threading layer to detect unsafe threads --- src/fast_array_utils/numba/__init__.py | 6 +++-- .../numba/_parallel_runtime.py | 22 +++++++++++++------ tests/test_numba.py | 15 +++++++++---- 3 files changed, 30 insertions(+), 13 deletions(-) diff --git a/src/fast_array_utils/numba/__init__.py b/src/fast_array_utils/numba/__init__.py index dde8e1f..dfd756b 100644 --- a/src/fast_array_utils/numba/__init__.py +++ b/src/fast_array_utils/numba/__init__.py @@ -79,8 +79,10 @@ def _threading_layer(layer_or_category: ThreadingLayer | TheadingCategory, /, pr def _is_on_unsafe_thread() -> bool: import threading - # We deem it unsafe if the caller is not the main thread, and therefore fall back to serial. - return threading.current_thread() is not threading.main_thread() and threading_layer() not in LAYERS["threadsafe"] + from ._parallel_runtime import _parallel_numba_runtime_layer + + # We deem it unsafe if the caller is not the main thread and the layer numba launches isn’t threadsafe, and therefore fall back to serial. + return threading.current_thread() is not threading.main_thread() and _parallel_numba_runtime_layer() not in LAYERS["threadsafe"] @overload diff --git a/src/fast_array_utils/numba/_parallel_runtime.py b/src/fast_array_utils/numba/_parallel_runtime.py index f30f1f1..db9e44c 100644 --- a/src/fast_array_utils/numba/_parallel_runtime.py +++ b/src/fast_array_utils/numba/_parallel_runtime.py @@ -7,7 +7,7 @@ import subprocess import sys from functools import cache -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast import numba @@ -18,7 +18,7 @@ from . import TheadingCategory, ThreadingLayer -__all__ = ["_needs_parallel_runtime_probe", "_parallel_numba_runtime_is_safe"] +__all__ = ["_needs_parallel_runtime_probe", "_parallel_numba_runtime_is_safe", "_parallel_numba_runtime_layer"] type _ParallelRuntimeProbeKey = tuple[str, ThreadingLayer | TheadingCategory, tuple[ThreadingLayer, ...], tuple[str, ...]] @@ -40,7 +40,7 @@ def _probe(values): values = np.arange(32, dtype=np.float64) assert _probe(values) == np.sum(values) -print({_PARALLEL_RUNTIME_PROBE_SENTINEL!r}) +print({_PARALLEL_RUNTIME_PROBE_SENTINEL!r}, numba.threading_layer()) """ @@ -87,7 +87,7 @@ def _build_parallel_runtime_probe_env(key: _ParallelRuntimeProbeKey | None = Non @cache -def _parallel_numba_runtime_is_safe_cached(key: _ParallelRuntimeProbeKey) -> bool: +def _parallel_numba_runtime_layer_cached(key: _ParallelRuntimeProbeKey) -> ThreadingLayer | None: try: # The probe command is built from `sys.executable` plus a generated script # that only imports modules from a fixed whitelist. @@ -100,9 +100,17 @@ def _parallel_numba_runtime_is_safe_cached(key: _ParallelRuntimeProbeKey) -> boo timeout=_PARALLEL_RUNTIME_PROBE_TIMEOUT, ) except Exception: # noqa: BLE001 - return False - return result.returncode == 0 and _PARALLEL_RUNTIME_PROBE_SENTINEL in result.stdout + return None + if result.returncode != 0: + return None + prefix = f"{_PARALLEL_RUNTIME_PROBE_SENTINEL} " + return next((cast("ThreadingLayer", line.removeprefix(prefix)) for line in result.stdout.splitlines() if line.startswith(prefix)), None) + + +def _parallel_numba_runtime_layer() -> ThreadingLayer | None: + """Get the threading layer numba actually launches, or `None` if parallel execution crashes.""" + return _parallel_numba_runtime_layer_cached(_parallel_runtime_probe_key()) def _parallel_numba_runtime_is_safe() -> bool: - return _parallel_numba_runtime_is_safe_cached(_parallel_runtime_probe_key()) + return _parallel_numba_runtime_layer() is not None diff --git a/tests/test_numba.py b/tests/test_numba.py index e95a0c5..d489982 100644 --- a/tests/test_numba.py +++ b/tests/test_numba.py @@ -40,7 +40,7 @@ def _sum_prange(values: NDArray[np.float64]) -> float: @pytest.fixture(autouse=True) def clear_probe_cache() -> None: - probe._parallel_numba_runtime_is_safe_cached.cache_clear() + probe._parallel_numba_runtime_layer_cached.cache_clear() fa_numba._threading_layer.cache_clear() @@ -90,14 +90,15 @@ def import_module(name: str, package: str | None = None) -> object: [ pytest.param(False, "workqueue", True, id="worker-unsafe"), pytest.param(False, "omp", False, id="worker-threadsafe"), + pytest.param(False, None, True, id="worker-probe-failed"), pytest.param(True, "workqueue", False, id="main-thread"), ], ) -def test_is_on_unsafe_thread(monkeypatch: pytest.MonkeyPatch, layer: fa_numba.ThreadingLayer, *, on_main: bool, expected: bool) -> None: +def test_is_on_unsafe_thread(monkeypatch: pytest.MonkeyPatch, layer: fa_numba.ThreadingLayer | None, *, on_main: bool, expected: bool) -> None: caller = threading.main_thread() if on_main else threading.Thread() monkeypatch.setattr(threading, "current_thread", lambda: caller) - monkeypatch.setattr(fa_numba, "threading_layer", lambda: layer) + monkeypatch.setattr(probe, "_parallel_numba_runtime_layer", lambda: layer) assert fa_numba._is_on_unsafe_thread() is expected @@ -222,10 +223,11 @@ def test_probe_result(monkeypatch: pytest.MonkeyPatch) -> None: def run(cmd: list[str], /, **kwargs: object) -> subprocess.CompletedProcess[str]: calls.append((cmd, kwargs)) - return subprocess.CompletedProcess(cmd, 0, stdout=f"{probe._PARALLEL_RUNTIME_PROBE_SENTINEL}\n", stderr="") + return subprocess.CompletedProcess(cmd, 0, stdout=f"{probe._PARALLEL_RUNTIME_PROBE_SENTINEL} tbb\n", stderr="") monkeypatch.setattr(probe.subprocess, "run", run) + assert probe._parallel_numba_runtime_layer() == "tbb" assert probe._parallel_numba_runtime_is_safe() is True assert probe._parallel_numba_runtime_is_safe() is True assert calls == [ @@ -242,6 +244,11 @@ def run(cmd: list[str], /, **kwargs: object) -> subprocess.CompletedProcess[str] ] +def test_probe_reports_launched_layer() -> None: + """Runs the real probe, which reports what numba launched instead of what we’d predict.""" + assert probe._parallel_numba_runtime_layer() in fa_numba.LAYERS["default"] + + @pytest.mark.parametrize( ("result", "error"), [ From 9c9a0a42cbbc12bd8a54658d47e65f9a2120302b Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Thu, 24 Sep 2026 14:59:19 +0200 Subject: [PATCH 2/3] =?UTF-8?q?refactor!:=20remove=20threading=5Flayer,=20?= =?UTF-8?q?which=20replicated=20numba=E2=80=99s=20layer=20selection?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/fast_array_utils/numba/__init__.py | 43 ++------------------------ tests/test_numba.py | 43 -------------------------- 2 files changed, 3 insertions(+), 83 deletions(-) diff --git a/src/fast_array_utils/numba/__init__.py b/src/fast_array_utils/numba/__init__.py index dfd756b..a5924df 100644 --- a/src/fast_array_utils/numba/__init__.py +++ b/src/fast_array_utils/numba/__init__.py @@ -13,16 +13,16 @@ import sys import warnings -from functools import cache, update_wrapper, wraps +from functools import update_wrapper, wraps from types import FunctionType from typing import TYPE_CHECKING, Literal, cast, overload if TYPE_CHECKING: - from collections.abc import Callable, Iterable + from collections.abc import Callable -__all__ = ["TheadingCategory", "ThreadingLayer", "njit", "threading_layer"] +__all__ = ["TheadingCategory", "ThreadingLayer", "njit"] type TheadingCategory = Literal["default", "safe", "threadsafe", "forksafe"] @@ -39,43 +39,6 @@ } -def threading_layer(layer_or_category: ThreadingLayer | TheadingCategory | None = None, /, priority: Iterable[ThreadingLayer] | None = None) -> ThreadingLayer: - """Get numba’s configured threading layer as specified in :ref:`numba-threading-layer`. - - ``layer_or_category`` defaults ``numba.config.THREADING_LAYER`` and ``priority`` to ``numba.config.THREADING_LAYER_PRIORITY``. - """ - import numba - - if layer_or_category is None: - layer_or_category = numba.config.THREADING_LAYER - if priority is None: - priority = numba.config.THREADING_LAYER_PRIORITY - - return _threading_layer(layer_or_category, tuple(priority)) - - -@cache -def _threading_layer(layer_or_category: ThreadingLayer | TheadingCategory, /, priority: Iterable[ThreadingLayer]) -> ThreadingLayer: - import importlib - - if (available := LAYERS.get(layer_or_category)) is None: # type: ignore[arg-type] # pragma: no cover - return cast("ThreadingLayer", layer_or_category) # given by direct name - - # given by layer type (safe, …) - for layer in priority: - if layer not in available: # pragma: no cover - continue - if layer != "workqueue": - try: # `importlib.util.find_spec` doesn’t work here - importlib.import_module(f"numba.np.ufunc.{layer}pool") - except ImportError: - continue - # the layer has been found - return layer - msg = f"No threading layer matching {layer_or_category!r} ({available=}, {priority=})" # pragma: no cover - raise ValueError(msg) # pragma: no cover - - def _is_on_unsafe_thread() -> bool: import threading diff --git a/tests/test_numba.py b/tests/test_numba.py index d489982..d13c6be 100644 --- a/tests/test_numba.py +++ b/tests/test_numba.py @@ -41,48 +41,6 @@ def _sum_prange(values: NDArray[np.float64]) -> float: @pytest.fixture(autouse=True) def clear_probe_cache() -> None: probe._parallel_numba_runtime_layer_cached.cache_clear() - fa_numba._threading_layer.cache_clear() - - -def test_threading_layer_uses_numba_config_defaults(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(numba.config, "THREADING_LAYER", "threadsafe") - monkeypatch.setattr(numba.config, "THREADING_LAYER_PRIORITY", ["omp", "tbb"]) - calls: list[tuple[fa_numba.ThreadingLayer | fa_numba.TheadingCategory, tuple[fa_numba.ThreadingLayer, ...]]] = [] - - def fake_threading_layer( - layer_or_category: fa_numba.ThreadingLayer | fa_numba.TheadingCategory, priority: tuple[fa_numba.ThreadingLayer, ...] - ) -> fa_numba.ThreadingLayer: - calls.append((layer_or_category, priority)) - return "omp" - - monkeypatch.setattr(fa_numba, "_threading_layer", fake_threading_layer) - - assert fa_numba.threading_layer() == "omp" - assert calls == [("threadsafe", ("omp", "tbb"))] - - -def test_threading_layer_resolves_available_backend(monkeypatch: pytest.MonkeyPatch) -> None: - original_import_module = importlib.import_module - calls: list[str] = [] - - def import_module(name: str, package: str | None = None) -> object: - calls.append(name) - if name.endswith("tbbpool"): - raise ImportError - if name.endswith("omppool"): - return object() - return original_import_module(name, package) - - monkeypatch.setattr(importlib, "import_module", import_module) - - assert fa_numba.threading_layer("threadsafe", ("workqueue", "tbb", "omp")) == "omp" - assert calls == ["numba.np.ufunc.tbbpool", "numba.np.ufunc.omppool"] - - fa_numba._threading_layer.cache_clear() - calls.clear() - - assert fa_numba.threading_layer("default", ("workqueue", "omp")) == "workqueue" - assert calls == [] @pytest.mark.parametrize( @@ -178,7 +136,6 @@ def test_probe_needed( def test_probe_check_is_lazy(monkeypatch: pytest.MonkeyPatch) -> None: _set_runtime(monkeypatch) - monkeypatch.setattr(fa_numba, "threading_layer", lambda: pytest.fail("threading_layer() should not be called")) original_import_module = importlib.import_module From 5d7e342546080a15ee2a8a3cc4fb69507c75cf79 Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Thu, 24 Sep 2026 16:12:32 +0200 Subject: [PATCH 3/3] refactor: inline _parallel_numba_runtime_is_safe --- src/fast_array_utils/numba/__init__.py | 4 ++-- src/fast_array_utils/numba/_parallel_runtime.py | 6 +----- tests/test_numba.py | 11 +++++------ 3 files changed, 8 insertions(+), 13 deletions(-) diff --git a/src/fast_array_utils/numba/__init__.py b/src/fast_array_utils/numba/__init__.py index a5924df..7c62a0e 100644 --- a/src/fast_array_utils/numba/__init__.py +++ b/src/fast_array_utils/numba/__init__.py @@ -63,7 +63,7 @@ def njit[**P, R](fn: Callable[P, R] | None = None, /) -> Callable[P, R] | Callab def decorator(f: Callable[P, R], /) -> Callable[P, R]: import numba - from ._parallel_runtime import _needs_parallel_runtime_probe, _parallel_numba_runtime_is_safe + from ._parallel_runtime import _needs_parallel_runtime_probe, _parallel_numba_runtime_layer assert isinstance(f, FunctionType) @@ -78,7 +78,7 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: msg = None if _is_on_unsafe_thread(): # pragma: no cover msg = f"Detected unsupported threading environment. Trying to run {f.__name__} in serial mode. In case of problems, install `tbb`." - elif _needs_parallel_runtime_probe() and not _parallel_numba_runtime_is_safe(): + elif _needs_parallel_runtime_probe() and _parallel_numba_runtime_layer() is None: msg = ( f"Detected an unsupported numba parallel runtime. Running {f.__name__} in serial mode as a workaround. " "Set `NUMBA_THREADING_LAYER=workqueue` or install `tbb` to avoid this fallback." diff --git a/src/fast_array_utils/numba/_parallel_runtime.py b/src/fast_array_utils/numba/_parallel_runtime.py index db9e44c..dce5cdd 100644 --- a/src/fast_array_utils/numba/_parallel_runtime.py +++ b/src/fast_array_utils/numba/_parallel_runtime.py @@ -18,7 +18,7 @@ from . import TheadingCategory, ThreadingLayer -__all__ = ["_needs_parallel_runtime_probe", "_parallel_numba_runtime_is_safe", "_parallel_numba_runtime_layer"] +__all__ = ["_needs_parallel_runtime_probe", "_parallel_numba_runtime_layer"] type _ParallelRuntimeProbeKey = tuple[str, ThreadingLayer | TheadingCategory, tuple[ThreadingLayer, ...], tuple[str, ...]] @@ -110,7 +110,3 @@ def _parallel_numba_runtime_layer_cached(key: _ParallelRuntimeProbeKey) -> Threa def _parallel_numba_runtime_layer() -> ThreadingLayer | None: """Get the threading layer numba actually launches, or `None` if parallel execution crashes.""" return _parallel_numba_runtime_layer_cached(_parallel_runtime_probe_key()) - - -def _parallel_numba_runtime_is_safe() -> bool: - return _parallel_numba_runtime_layer() is not None diff --git a/tests/test_numba.py b/tests/test_numba.py index d13c6be..6a04bad 100644 --- a/tests/test_numba.py +++ b/tests/test_numba.py @@ -185,8 +185,7 @@ def run(cmd: list[str], /, **kwargs: object) -> subprocess.CompletedProcess[str] monkeypatch.setattr(probe.subprocess, "run", run) assert probe._parallel_numba_runtime_layer() == "tbb" - assert probe._parallel_numba_runtime_is_safe() is True - assert probe._parallel_numba_runtime_is_safe() is True + assert probe._parallel_numba_runtime_layer() == "tbb" # cached assert calls == [ ( [probe.sys.executable, "-c", probe._parallel_runtime_probe_code(("torch",))], @@ -230,7 +229,7 @@ def run(_cmd: list[str], /, **_kwargs: object) -> subprocess.CompletedProcess[st monkeypatch.setattr(probe.subprocess, "run", run) - assert probe._parallel_numba_runtime_is_safe() is False + assert probe._parallel_numba_runtime_layer() is None @pytest.mark.parametrize( @@ -260,9 +259,9 @@ def test_njit_chooses_version( else: monkeypatch.setattr(probe, "_needs_parallel_runtime_probe", lambda: needs_probe) if probe_safe is None: - monkeypatch.setattr(probe, "_parallel_numba_runtime_is_safe", lambda: pytest.fail("probe should not run")) + monkeypatch.setattr(probe, "_parallel_numba_runtime_layer", lambda: pytest.fail("probe should not run")) else: - monkeypatch.setattr(probe, "_parallel_numba_runtime_is_safe", lambda: probe_safe) + monkeypatch.setattr(probe, "_parallel_numba_runtime_layer", lambda: "tbb" if probe_safe else None) wrapped = fa_numba.njit(_return_true) @@ -281,7 +280,7 @@ def test_serial_fallback(monkeypatch: pytest.MonkeyPatch) -> None: values = np.arange(10, dtype=np.float64) monkeypatch.setattr(fa_numba, "_is_on_unsafe_thread", lambda: False) monkeypatch.setattr(probe, "_needs_parallel_runtime_probe", lambda: True) - monkeypatch.setattr(probe, "_parallel_numba_runtime_is_safe", lambda: False) + monkeypatch.setattr(probe, "_parallel_numba_runtime_layer", lambda: None) wrapped = fa_numba.njit(_sum_prange) with pytest.warns(UserWarning, match="unsupported numba parallel runtime"):