fix: use parallelized numba functions if possible

Author: ilan-gold

pyproject.toml

@@ -177,12 +177,6 @@ doctest_subpackage_requires = [
   "src/fast_array_utils/_plugins/dask.py = dask",
   "src/fast_array_utils/_plugins/numba_sparse.py = numba;scipy",
 ]
-filterwarnings = [
-  "error",
-  # codspeed seems to break this dtype added by h5py
-  "ignore:.*numpy[.]longdouble:UserWarning",
-  "ignore:FNV hashing is not implemented in Numba:UserWarning",
-]
 markers = [
   "benchmark: marks tests as benchmark (to run with `--codspeed`)",
 ]

src/fast_array_utils/conv/scipy/_to_dense.py

@@ -5,6 +5,8 @@
 
 import numba
 
+from fast_array_utils.utils import njit
+
 
 if TYPE_CHECKING:
     from typing import Any
@@ -18,14 +20,14 @@
 __all__ = ["_to_dense_csc_numba", "_to_dense_csr_numba"]
 
 
-@numba.njit(cache=True)
+@njit
 def _to_dense_csc_numba(x: CSBase, out: NDArray[np.number[Any]]) -> None:
     for c in numba.prange(out.shape[1]):
         for i in range(x.indptr[c], x.indptr[c + 1]):
             out[x.indices[i], c] = x.data[i]
 
 
-@numba.njit(cache=True)
+@njit
 def _to_dense_csr_numba(x: CSBase, out: NDArray[np.number[Any]]) -> None:
     for r in numba.prange(out.shape[0]):
         for i in range(x.indptr[r], x.indptr[r + 1]):

src/fast_array_utils/stats/_is_constant.py

@@ -7,6 +7,8 @@
 import numba
 import numpy as np
 
+from fast_array_utils.utils import njit
+
 from .. import types
 
 
@@ -64,7 +66,7 @@ def _is_constant_cs(a: types.CSBase, /, *, axis: Literal[0, 1] | None = None) ->
     return _is_constant_cs_major(a, shape)
 
 
-@numba.njit(cache=True)
+@njit
 def _is_constant_cs_major(a: types.CSBase, shape: tuple[int, int]) -> NDArray[np.bool]:
     n = len(a.indptr) - 1
     result = np.ones(n, dtype=np.bool)

src/fast_array_utils/stats/_mean_var.py

@@ -6,6 +6,8 @@
 import numba
 import numpy as np
 
+from fast_array_utils.utils import njit
+
 from .. import types
 from ._power import power
 
@@ -79,7 +81,7 @@ def _sparse_mean_var(mtx: types.CSBase, /, *, axis: Literal[0, 1]) -> tuple[NDAr
     )
 
 
-@numba.njit(cache=True)
+@njit
 def sparse_mean_var_minor_axis(
     x: types.CSBase,
     *,
@@ -109,7 +111,7 @@ def sparse_mean_var_minor_axis(
     return means, variances
 
 
-@numba.njit(cache=True)
+@njit
 def sparse_mean_var_major_axis(
     x: types.CSBase,
     *,

src/fast_array_utils/utils.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+import sys
+import warnings
+from functools import cache, wraps
+from typing import TYPE_CHECKING, Literal, cast, overload
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+type LayerType = Literal["default", "safe", "threadsafe", "forksafe"]
+type Layer = Literal["tbb", "omp", "workqueue"]
+
+
+LAYERS: dict[LayerType, set[Layer]] = {
+    "default": {"tbb", "omp", "workqueue"},
+    "safe": {"tbb"},
+    "threadsafe": {"tbb", "omp"},
+    "forksafe": {"tbb", "workqueue", *(() if sys.platform == "linux" else {"omp"})},
+}
+
+
+@cache
+def _numba_threading_layer() -> Layer:
+    """Get numba’s threading layer.
+
+    This function implements the algorithm as described in
+    <https://numba.readthedocs.io/en/stable/user/threading-layer.html>
+    """
+    import importlib
+
+    import numba
+
+    if (available := LAYERS.get(numba.config.THREADING_LAYER)) is None:
+        # given by direct name
+        return numba.config.THREADING_LAYER
+
+    # given by layer type (safe, …)
+    for layer in cast("list[Layer]", numba.config.THREADING_LAYER_PRIORITY):
+        if layer not in available:
+            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 loadable threading layer: {numba.config.THREADING_LAYER=}  ({available=}, {numba.config.THREADING_LAYER_PRIORITY=})"
+    raise ValueError(msg)
+
+
+def _is_in_unsafe_thread_pool() -> bool:
+    import threading
+
+    current_thread = threading.current_thread()
+    # ThreadPoolExecutor threads typically have names like 'ThreadPoolExecutor-0_1'
+    return current_thread.name.startswith("ThreadPoolExecutor") and _numba_threading_layer() not in LAYERS["threadsafe"]
+
+
+@overload
+def njit[**P, R](fn: Callable[P, R], /) -> Callable[P, R]: ...
+@overload
+def njit[**P, R]() -> Callable[[Callable[P, R]], Callable[P, R]]: ...
+def njit[**P, R](fn: Callable[P, R] | None = None, /) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]:
+    """Jit-compile a function using numba.
+
+    On call, this function dispatches to a parallel or sequential numba function,
+    depending on if it has been called from a thread pool.
+
+    See <https://github.com/numbagg/numbagg/pull/201/files#r1409374809>
+    """
+
+    def decorator(f: Callable[P, R], /) -> Callable[P, R]:
+        import numba
+
+        fns: dict[bool, Callable[P, R]] = {parallel: numba.njit(f, cache=True, parallel=parallel) for parallel in (True, False)}
+
+        @wraps(f)
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+            parallel = not _is_in_unsafe_thread_pool()
+            if not parallel:
+                msg = f"Detected unsupported threading environment. Trying to run {f.__name__} in serial mode. In case of problems, install `tbb`."
+                warnings.warn(msg, UserWarning, stacklevel=2)
+            return fns[parallel](*args, **kwargs)
+
+        return wrapper
+
+    return decorator if fn is None else decorator(fn)

src/testing/fast_array_utils/pytest.py

@@ -89,6 +89,23 @@ def _skip_if_unimportable(array_type: ArrayType) -> pytest.MarkDecorator:
 SUPPORTED_TYPE_PARAMS = [pytest.param(t, id=str(t), marks=_skip_if_unimportable(t)) for t in SUPPORTED_TYPES]
 
 
+@pytest.fixture(autouse=True)
+def dask_single_threaded() -> Generator[None]:
+    """Switch to a single-threaded scheduler for tests since numba crashes otherwise."""
+    if not find_spec("dask"):
+        yield
+        return
+
+    import dask.config
+
+    prev_scheduler = dask.config.get("scheduler", "threads")
+    dask.config.set(scheduler="single-threaded")
+    try:
+        yield
+    finally:
+        dask.config.set(scheduler=prev_scheduler)
+
+
 @pytest.fixture(scope="session", params=SUPPORTED_TYPE_PARAMS)
 def array_type(request: pytest.FixtureRequest) -> ArrayType:
     """Fixture for a supported :class:`~testing.fast_array_utils.ArrayType`.