fix!: use the probed threading layer to detect unsafe threads

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,48 +39,13 @@
 }
 
 
-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
 
-    # 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
@@ -98,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)
 
@@ -113,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."

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_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,13 @@ 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_is_safe() -> bool:
-    return _parallel_numba_runtime_is_safe_cached(_parallel_runtime_probe_key())
+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())

tests/test_numba.py

@@ -40,64 +40,23 @@ 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()
-    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 == []
+    probe._parallel_numba_runtime_layer_cached.cache_clear()
 
 
 @pytest.mark.parametrize(
     ("on_main", "layer", "expected"),
     [
         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
 
@@ -177,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
 
@@ -222,12 +180,12 @@ 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_is_safe() is True
-    assert probe._parallel_numba_runtime_is_safe() is True
+    assert probe._parallel_numba_runtime_layer() == "tbb"
+    assert probe._parallel_numba_runtime_layer() == "tbb"  # cached
     assert calls == [
         (
             [probe.sys.executable, "-c", probe._parallel_runtime_probe_code(("torch",))],
@@ -242,6 +200,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"),
     [
@@ -266,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(
@@ -296,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)
 
@@ -317,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"):