From 5dddda47de18d560902cf7a8088092954c23c4e8 Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Fri, 25 Sep 2026 13:13:49 +0200 Subject: [PATCH 1/3] feat: add nogil option to njit (default True) --- src/fast_array_utils/numba/__init__.py | 15 +++++++++++---- tests/test_numba.py | 10 ++++++++-- 2 files changed, 19 insertions(+), 6 deletions(-) diff --git a/src/fast_array_utils/numba/__init__.py b/src/fast_array_utils/numba/__init__.py index 7c62a0e..fd5c1dd 100644 --- a/src/fast_array_utils/numba/__init__.py +++ b/src/fast_array_utils/numba/__init__.py @@ -49,14 +49,19 @@ def _is_on_unsafe_thread() -> bool: @overload -def njit[**P, R](fn: Callable[P, R], /) -> Callable[P, R]: ... +def njit[**P, R](fn: Callable[P, R], /, *, nogil: bool = True) -> 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]]: +def njit[**P, R](*, nogil: bool = True) -> Callable[[Callable[P, R]], Callable[P, R]]: ... +def njit[**P, R](fn: Callable[P, R] | None = None, /, *, nogil: bool = True) -> 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 serial numba function, depending on the current threading environment. + + Parameters + ---------- + nogil + Release the GIL while the compiled function runs (see :func:`numba.jit`). """ # See https://github.com/numbagg/numbagg/pull/201/files#r1409374809 @@ -69,7 +74,9 @@ def decorator(f: Callable[P, R], /) -> Callable[P, R]: # use distinct names so numba doesn’t reuse the wrong version’s cache fns: dict[bool, Callable[P, R]] = { - parallel: numba.njit(_copy_function(f, __qualname__=f"{f.__qualname__}-{'parallel' if parallel else 'serial'}"), cache=True, parallel=parallel) + parallel: numba.njit( + _copy_function(f, __qualname__=f"{f.__qualname__}-{'parallel' if parallel else 'serial'}"), cache=True, parallel=parallel, nogil=nogil + ) for parallel in (True, False) } diff --git a/tests/test_numba.py b/tests/test_numba.py index 6a04bad..b71ce2f 100644 --- a/tests/test_numba.py +++ b/tests/test_numba.py @@ -83,9 +83,10 @@ def _set_runtime( monkeypatch.setattr(probe, "LAYERS", layers) -def _install_fake_njit(monkeypatch: pytest.MonkeyPatch, calls: list[bool]) -> None: - def fake_njit(_fn: object, /, *, cache: bool, parallel: bool) -> Callable[..., bool]: +def _install_fake_njit(monkeypatch: pytest.MonkeyPatch, calls: list[bool], *, expected_nogil: bool = True) -> None: + def fake_njit(_fn: object, /, *, cache: bool, parallel: bool, nogil: bool) -> Callable[..., bool]: assert cache is True + assert nogil is expected_nogil def compiled(*_args: object, **_kwargs: object) -> bool: calls.append(parallel) @@ -287,3 +288,8 @@ def test_serial_fallback(monkeypatch: pytest.MonkeyPatch) -> None: result = wrapped(values) assert result == pytest.approx(np.sum(values)) + + +def test_njit_nogil(monkeypatch: pytest.MonkeyPatch) -> None: + _install_fake_njit(monkeypatch, [], expected_nogil=False) + fa_numba.njit(nogil=False)(_return_true) From d82cf7c5bdd048e162172d6dfbf116ae671fba42 Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Fri, 25 Sep 2026 13:13:49 +0200 Subject: [PATCH 2/3] test: add threaded mean_var benchmark --- src/fast_array_utils/numba/__init__.py | 2 ++ tests/test_stats.py | 15 +++++++++++++++ 2 files changed, 17 insertions(+) diff --git a/src/fast_array_utils/numba/__init__.py b/src/fast_array_utils/numba/__init__.py index fd5c1dd..e654a93 100644 --- a/src/fast_array_utils/numba/__init__.py +++ b/src/fast_array_utils/numba/__init__.py @@ -62,6 +62,8 @@ def njit[**P, R](fn: Callable[P, R] | None = None, /, *, nogil: bool = True) -> ---------- nogil Release the GIL while the compiled function runs (see :func:`numba.jit`). + numba’s on-disk cache doesn’t distinguish this option, + so wrapping the same function with both values reuses whichever was compiled first. """ # See https://github.com/numbagg/numbagg/pull/201/files#r1409374809 diff --git a/tests/test_stats.py b/tests/test_stats.py index 0860334..4d5afc6 100644 --- a/tests/test_stats.py +++ b/tests/test_stats.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: MPL-2.0 from __future__ import annotations +from concurrent.futures import ThreadPoolExecutor +from functools import partial from importlib.metadata import version from importlib.util import find_spec from pathlib import Path @@ -409,3 +411,16 @@ def test_stats_benchmark( @benchmark def call() -> None: func(arr, axis=axis) + + +@pytest.mark.benchmark +@pytest.mark.array_type(select=Flags.Sparse, skip=Flags.Matrix | Flags.Dask | Flags.Disk | Flags.Gpu) +def test_mean_var_threaded_benchmark(benchmark: BenchmarkFixture, array_type: ArrayType[CpuArray, None]) -> None: + """Calls from worker threads run serially, so they only overlap if numba releases the GIL.""" + mats = [array_type.random((20_000, 5_000), density=0.1, dtype=np.float64) for _ in range(4)] + stats.mean_var(mats[0], axis=0) # warmup: numba compile + + @benchmark + def call() -> None: + with ThreadPoolExecutor(len(mats)) as pool: + list(pool.map(partial(stats.mean_var, axis=0), mats)) From 1d69e04bcadbe3adf938b2872e4d2299395ff162 Mon Sep 17 00:00:00 2001 From: "Philipp A." Date: Fri, 25 Sep 2026 13:38:25 +0200 Subject: [PATCH 3/3] docs: reduce --- src/fast_array_utils/numba/__init__.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/fast_array_utils/numba/__init__.py b/src/fast_array_utils/numba/__init__.py index e654a93..fd5c1dd 100644 --- a/src/fast_array_utils/numba/__init__.py +++ b/src/fast_array_utils/numba/__init__.py @@ -62,8 +62,6 @@ def njit[**P, R](fn: Callable[P, R] | None = None, /, *, nogil: bool = True) -> ---------- nogil Release the GIL while the compiled function runs (see :func:`numba.jit`). - numba’s on-disk cache doesn’t distinguish this option, - so wrapping the same function with both values reuses whichever was compiled first. """ # See https://github.com/numbagg/numbagg/pull/201/files#r1409374809