Backport PR scverse#4015 on branch 1.12.x (fix: use fast-array-utils’ threadsafe njit) (scverse#4022)

Author: flying-sheep

.github/workflows/ci.yml

@@ -27,7 +27,7 @@ jobs:
         with: { enable-cache: false }
       - id: get-envs
         run: |
-          ENVS_JSON=$(NO_COLOR=1 uvx '--with=virtualenv<21' hatch env show --json | jq -c 'to_entries
+          ENVS_JSON=$(NO_COLOR=1 uvx hatch env show --json | jq -c 'to_entries
             | map(
               select(.key | startswith("hatch-test"))
               | {
@@ -63,7 +63,7 @@ jobs:
       - name: Install dependencies
         run: |
           echo "::group::Install hatch"
-          uv tool install hatch '--with=virtualenv<21'
+          uv tool install hatch
           echo "::endgroup::"
           echo "::group::Create environment"
           hatch -v env create ${{ matrix.env.name }}

.readthedocs.yml

@@ -16,8 +16,8 @@ build:
       - asdf global uv latest
     pre_build:
       # run towncrier to preview the next version’s release notes
-      - ( find docs/release-notes -regex '[^.]+[.][^.]+.md' | grep -q . ) && uvx "--with=virtualenv<21" hatch run towncrier build --keep || true
+      - ( find docs/release-notes -regex '[^.]+[.][^.]+.md' | grep -q . ) && uvx hatch run towncrier build --keep || true
     build:
       html:
-        - uvx "--with=virtualenv<21" hatch run docs:build
+        - uvx hatch run docs:build
         - mv docs/_build $READTHEDOCS_OUTPUT

docs/release-notes/4015.fix.md

@@ -0,0 +1 @@
+Fix crashes when running numba in dask by using the threadsafe {func}`fast_array_utils.numba.njit` {smaller}`P Angerer`

pyproject.toml

@@ -53,7 +53,7 @@ dynamic = [ "version" ]
 dependencies = [
     "anndata>=0.10.8",
     "certifi",
-    "fast-array-utils[accel,sparse]>=1.2.1",
+    "fast-array-utils[accel,sparse]>=1.4",
     "h5py>=3.11",
     "joblib",
     "legacy-api-wrap>=1.5",                     # for positional API deprecations
@@ -215,8 +215,8 @@ lint.pylint.max-positional-args = 5
 
 [tool.ruff.lint.flake8-tidy-imports.banned-api]
 "legacy_api_wrap.legacy_api".msg = "Use scanpy._compat.old_positionals instead"
-"numba.jit".msg = "Use `scanpy._compat.njit` instead"
-"numba.njit".msg = "Use `scanpy._compat.njit` instead"
+"numba.jit".msg = "Use `fast_array_utils.numba.njit` instead"
+"numba.njit".msg = "Use `fast_array_utils.numba.njit` instead"
 "numpy.bool_".msg = "Use `np.bool` instead"
 "pandas.api.types.is_categorical_dtype".msg = "Use isinstance(s.dtype, CategoricalDtype) instead"
 "pandas.value_counts".msg = "Use pd.Series(a).value_counts() instead"
@@ -291,7 +291,7 @@ run.source_pkgs = [ "scanpy" ]
 paths.source = [ "src", "**/site-packages" ]
 report.exclude_also = [
     # https://github.com/numba/numba/issues/4268
-    "@(numba\\.|nb\\.)?njit.*",
+    "@([\\w.]+.)?njit.*",
     "@deprecated.*",
     "if __name__ == .__main__.:",
     "if TYPE_CHECKING:",

src/scanpy/_compat.py

@@ -2,17 +2,16 @@
 
 import sys
 import warnings
-from functools import cache, partial, wraps
+from functools import cache, partial
 from importlib.util import find_spec
 from pathlib import Path
-from typing import TYPE_CHECKING, Literal, cast, overload
+from typing import TYPE_CHECKING
 
 import legacy_api_wrap
 from packaging.version import Version
 from scipy import sparse
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
     from importlib.metadata import PackageMetadata
 
 
@@ -22,10 +21,8 @@
     "CSRBase",
     "DaskArray",
     "SpBase",
-    "_numba_threading_layer",
     "deprecated",
     "fullname",
-    "njit",
     "old_positionals",
     "pkg_metadata",
     "pkg_version",
@@ -123,99 +120,3 @@ def warn(
     warnings.warn(  # noqa: TID251
         message, category, source=source, skip_file_prefixes=skip_file_prefixes
     )
-
-
-@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)  # 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`."
-                )
-                warn(msg, UserWarning)
-            return fns[parallel](*args, **kwargs)
-
-        return wrapper
-
-    return decorator if fn is None else decorator(fn)
-
-
-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"})},
-}
-
-
-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/experimental/pp/_highly_variable_genes.py

@@ -8,10 +8,11 @@
 import numpy as np
 import pandas as pd
 from anndata import AnnData
+from fast_array_utils.numba import njit
 from fast_array_utils.stats import mean_var
 
 from ... import logging as logg
-from ..._compat import CSBase, njit, warn
+from ..._compat import CSBase, warn
 from ..._settings import Verbosity, settings
 from ..._utils import _doc_params, check_nonnegative_integers, view_to_actual
 from ...experimental._docs import (

src/scanpy/metrics/_gearys_c.py

@@ -7,8 +7,9 @@
 
 import numba
 import numpy as np
+from fast_array_utils.numba import njit
 
-from .._compat import CSRBase, njit
+from .._compat import CSRBase
 from .._utils import _doc_params
 from ..get import _get_obs_rep
 from ..neighbors._doc import doc_neighbors_key

src/scanpy/metrics/_morans_i.py

@@ -7,8 +7,9 @@
 
 import numba
 import numpy as np
+from fast_array_utils.numba import njit
 
-from .._compat import CSRBase, njit
+from .._compat import CSRBase
 from .._utils import _doc_params
 from ..get import _get_obs_rep
 from ..neighbors._doc import doc_neighbors_key

src/scanpy/preprocessing/_normalization.py

@@ -6,9 +6,10 @@
 import numba
 import numpy as np
 from fast_array_utils import stats
+from fast_array_utils.numba import njit
 
 from .. import logging as logg
-from .._compat import CSBase, CSCBase, CSRBase, DaskArray, njit, old_positionals, warn
+from .._compat import CSBase, CSCBase, CSRBase, DaskArray, old_positionals, warn
 from .._utils import axis_mul_or_truediv, dematrix, view_to_actual
 from ..get import _get_obs_rep, _set_obs_rep
 

src/scanpy/preprocessing/_qc.py

@@ -7,12 +7,13 @@
 import numpy as np
 import pandas as pd
 from fast_array_utils import stats
+from fast_array_utils.numba import njit
 from scipy import sparse
 
 from scanpy.get import _get_obs_rep
 from scanpy.preprocessing._distributed import materialize_as_ndarray
 
-from .._compat import CSBase, CSRBase, DaskArray, njit, warn
+from .._compat import CSBase, CSRBase, DaskArray, warn
 from .._utils import _doc_params, axis_nnz
 from ._docs import (
     doc_adata_basic,