move all njit calls into a decorator (scverse#3335)

Author: flying-sheep

docs/release-notes/3335.feature.md

@@ -0,0 +1 @@
+Run numba functions single-threaded when called from inside of a ThreadPool {smaller}`P Angerer`

pyproject.toml

@@ -262,6 +262,8 @@ required-imports = ["from __future__ import annotations"]
 "pandas.value_counts".msg = "Use pd.Series(a).value_counts() instead"
 "legacy_api_wrap.legacy_api".msg = "Use scanpy._compat.old_positionals instead"
 "numpy.bool".msg = "Use `np.bool_` instead for numpy>=1.24<2 compatibility"
+"numba.jit".msg = "Use `scanpy._compat.njit` instead"
+"numba.njit".msg = "Use `scanpy._compat.njit` instead"
 [tool.ruff.lint.flake8-type-checking]
 exempt-modules = []
 strict = true

src/scanpy/_compat.py

@@ -1,17 +1,23 @@
 from __future__ import annotations
 
+import os
 import sys
+import warnings
 from dataclasses import dataclass, field
-from functools import cache, partial
+from functools import cache, partial, wraps
 from importlib.util import find_spec
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Literal, ParamSpec, TypeVar, cast, overload
 
 from packaging.version import Version
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
     from importlib.metadata import PackageMetadata
 
+P = ParamSpec("P")
+R = TypeVar("R")
+
 
 if TYPE_CHECKING:
     # type checkers are confused and can only see …core.Array
@@ -90,3 +96,101 @@ def pkg_version(package: str) -> Version:
     # but this code makes it possible to run scanpy without it.
     def old_positionals(*old_positionals: str):
         return lambda func: func
+
+
+@overload
+def njit(fn: Callable[P, R], /) -> Callable[P, R]: ...
+@overload
+def njit() -> Callable[[Callable[P, R]], Callable[P, R]]: ...
+def njit(
+    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)  # noqa: TID251
+            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 = (
+                    "Detected unsupported threading environment. "
+                    f"Trying to run {f.__name__} in serial mode. "
+                    "In case of problems, install `tbb`."
+                )
+                warnings.warn(msg, stacklevel=2)
+            return fns[parallel](*args, **kwargs)
+
+        return wrapper
+
+    return decorator if fn is None else decorator(fn)
+
+
+LayerType = Literal["default", "safe", "threadsafe", "forksafe"]
+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"})},
+}
+
+
+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"]
+    )
+
+
+@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=} "
+        f" ({available=}, {numba.config.THREADING_LAYER_PRIORITY=})"
+    )
+    raise ValueError(msg)

src/scanpy/_utils/compute/is_constant.py

@@ -5,11 +5,11 @@
 from numbers import Integral
 from typing import TYPE_CHECKING, TypeVar, overload
 
+import numba
 import numpy as np
-from numba import njit
 from scipy import sparse
 
-from ..._compat import DaskArray
+from ..._compat import DaskArray, njit
 
 if TYPE_CHECKING:
     from typing import Literal
@@ -103,22 +103,21 @@ def _(
         else:
             return (a.data == 0).all()
     if axis == 1:
-        return _is_constant_csr_rows(a.data, a.indices, a.indptr, a.shape)
+        return _is_constant_csr_rows(a.data, a.indptr, a.shape)
     elif axis == 0:
         a = a.T.tocsr()
-        return _is_constant_csr_rows(a.data, a.indices, a.indptr, a.shape)
+        return _is_constant_csr_rows(a.data, a.indptr, a.shape)
 
 
 @njit
 def _is_constant_csr_rows(
     data: NDArray[np.number],
-    indices: NDArray[np.integer],
     indptr: NDArray[np.integer],
     shape: tuple[int, int],
-):
+) -> NDArray[np.bool_]:
     n = len(indptr) - 1
     result = np.ones(n, dtype=np.bool_)
-    for i in range(n):
+    for i in numba.prange(n):
         start = indptr[i]
         stop = indptr[i + 1]
         val = data[start] if stop - start == shape[1] else 0
@@ -139,10 +138,10 @@ def _(
         else:
             return (a.data == 0).all()
     if axis == 0:
-        return _is_constant_csr_rows(a.data, a.indices, a.indptr, a.shape[::-1])
+        return _is_constant_csr_rows(a.data, a.indptr, a.shape[::-1])
     elif axis == 1:
         a = a.T.tocsc()
-        return _is_constant_csr_rows(a.data, a.indices, a.indptr, a.shape[::-1])
+        return _is_constant_csr_rows(a.data, a.indptr, a.shape[::-1])
 
 
 @is_constant.register(DaskArray)
@@ -151,4 +150,8 @@ def _(a: DaskArray, axis: Literal[0, 1] | None = None) -> bool | NDArray[np.bool
         v = a[tuple(0 for _ in range(a.ndim))].compute()
         return (a == v).all()
     # TODO: use overlapping blocks and reduction instead of `drop_axis`
-    return a.map_blocks(partial(is_constant, axis=axis), drop_axis=axis)
+    return a.map_blocks(
+        partial(is_constant, axis=axis),
+        drop_axis=axis,
+        meta=np.array([], dtype=a.dtype),
+    )

src/scanpy/experimental/pp/_highly_variable_genes.py

@@ -12,6 +12,7 @@
 from anndata import AnnData
 
 from scanpy import logging as logg
+from scanpy._compat import njit
 from scanpy._settings import Verbosity, settings
 from scanpy._utils import _doc_params, check_nonnegative_integers, view_to_actual
 from scanpy.experimental._docs import (
@@ -32,7 +33,7 @@
     from numpy.typing import NDArray
 
 
-@nb.njit(parallel=True)
+@njit
 def _calculate_res_sparse(
     indptr: NDArray[np.integer],
     index: NDArray[np.integer],
@@ -92,7 +93,7 @@ def clac_clipped_res_sparse(gene: int, cell: int, value: np.float64) -> np.float
     return residuals
 
 
-@nb.njit(parallel=True)
+@njit
 def _calculate_res_dense(
     matrix,
     *,

src/scanpy/metrics/_gearys_c.py

@@ -9,7 +9,7 @@
 import numpy as np
 from scipy import sparse
 
-from .._compat import fullname
+from .._compat import fullname, njit
 from ..get import _get_obs_rep
 from ._common import _check_vals, _resolve_vals
 
@@ -136,7 +136,6 @@ def gearys_c(
 #   tests to fail.
 
 
-@numba.njit(cache=True, parallel=True)
 def _gearys_c_vec(
     data: np.ndarray,
     indices: np.ndarray,
@@ -147,7 +146,7 @@ def _gearys_c_vec(
     return _gearys_c_vec_W(data, indices, indptr, x, W)
 
 
-@numba.njit(cache=True, parallel=True)
+@njit
 def _gearys_c_vec_W(
     data: np.ndarray,
     indices: np.ndarray,
@@ -182,7 +181,7 @@ def _gearys_c_vec_W(
 # https://github.com/numba/numba/issues/6774#issuecomment-788789663
 
 
-@numba.njit(cache=True)
+@numba.njit(cache=True, parallel=False)  # noqa: TID251
 def _gearys_c_inner_sparse_x_densevec(
     g_data: np.ndarray,
     g_indices: np.ndarray,
@@ -203,7 +202,7 @@ def _gearys_c_inner_sparse_x_densevec(
     return numer / denom
 
 
-@numba.njit(cache=True)
+@numba.njit(cache=True, parallel=False)  # noqa: TID251
 def _gearys_c_inner_sparse_x_sparsevec(  # noqa: PLR0917
     g_data: np.ndarray,
     g_indices: np.ndarray,
@@ -239,7 +238,7 @@ def _gearys_c_inner_sparse_x_sparsevec(  # noqa: PLR0917
     return numer / denom
 
 
-@numba.njit(cache=True, parallel=True)
+@njit
 def _gearys_c_mtx(
     g_data: np.ndarray,
     g_indices: np.ndarray,
@@ -256,7 +255,7 @@ def _gearys_c_mtx(
     return out
 
 
-@numba.njit(cache=True, parallel=True)
+@njit
 def _gearys_c_mtx_csr(  # noqa: PLR0917
     g_data: np.ndarray,
     g_indices: np.ndarray,

src/scanpy/metrics/_morans_i.py

@@ -9,7 +9,7 @@
 import numpy as np
 from scipy import sparse
 
-from .._compat import fullname
+from .._compat import fullname, njit
 from ..get import _get_obs_rep
 from ._common import _check_vals, _resolve_vals
 
@@ -126,7 +126,7 @@ def morans_i(
 # This is done in a very similar way to gearys_c. See notes there for details.
 
 
-@numba.njit(cache=True, parallel=True)
+@njit
 def _morans_i_vec(
     g_data: np.ndarray,
     g_indices: np.ndarray,
@@ -137,7 +137,7 @@ def _morans_i_vec(
     return _morans_i_vec_W(g_data, g_indices, g_indptr, x, W)
 
 
-@numba.njit(cache=True)
+@numba.njit(cache=True, parallel=False)  # noqa: TID251
 def _morans_i_vec_W(
     g_data: np.ndarray,
     g_indices: np.ndarray,
@@ -159,7 +159,7 @@ def _morans_i_vec_W(
     return len(x) / W * inum / z2ss
 
 
-@numba.njit(cache=True)
+@numba.njit(cache=True, parallel=False)  # noqa: TID251
 def _morans_i_vec_W_sparse(  # noqa: PLR0917
     g_data: np.ndarray,
     g_indices: np.ndarray,
@@ -174,7 +174,7 @@ def _morans_i_vec_W_sparse(  # noqa: PLR0917
     return _morans_i_vec_W(g_data, g_indices, g_indptr, x, W)
 
 
-@numba.njit(cache=True, parallel=True)
+@njit
 def _morans_i_mtx(
     g_data: np.ndarray,
     g_indices: np.ndarray,
@@ -191,7 +191,7 @@ def _morans_i_mtx(
     return out
 
 
-@numba.njit(cache=True, parallel=True)
+@njit
 def _morans_i_mtx_csr(  # noqa: PLR0917
     g_data: np.ndarray,
     g_indices: np.ndarray,

src/scanpy/preprocessing/_highly_variable_genes.py

@@ -200,7 +200,8 @@ def _highly_variable_genes_seurat_v3(
         return df
 
 
-@numba.njit(cache=True)
+# parallel=False needed for accuracy
+@numba.njit(cache=True, parallel=False)  # noqa: TID251
 def _sum_and_sum_squares_clipped(
     indices: NDArray[np.integer],
     data: NDArray[np.floating],
@@ -211,7 +212,7 @@ def _sum_and_sum_squares_clipped(
 ) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
     squared_batch_counts_sum = np.zeros(n_cols, dtype=np.float64)
     batch_counts_sum = np.zeros(n_cols, dtype=np.float64)
-    for i in range(nnz):
+    for i in numba.prange(nnz):
         idx = indices[i]
         element = min(np.float64(data[i]), clip_val[idx])
         squared_batch_counts_sum[idx] += element**2

src/scanpy/preprocessing/_qc.py

@@ -12,7 +12,7 @@
 from scanpy.preprocessing._distributed import materialize_as_ndarray
 from scanpy.preprocessing._utils import _get_mean_var
 
-from .._compat import DaskArray
+from .._compat import DaskArray, njit
 from .._utils import _doc_params, axis_nnz, axis_sum
 from ._docs import (
     doc_adata_basic,
@@ -445,7 +445,7 @@ def _(mtx: spmatrix, ns: Collection[int]) -> DaskArray:
     return top_segment_proportions_sparse_csr(mtx.data, mtx.indptr, np.array(ns))
 
 
-@numba.njit(cache=True, parallel=True)
+@njit
 def top_segment_proportions_sparse_csr(data, indptr, ns):
     # work around https://github.com/numba/numba/issues/5056
     indptr = indptr.astype(np.int64)