Merge pull request #1 from scverse/fix-rgg-n_jobs

move to utils, slimmer tests

Author: JhonatanFelix

src/scanpy/_utils/__init__.py

@@ -34,6 +34,7 @@
 
 from .. import logging as logg
 from .._compat import CSBase, DaskArray, _CSArray, pkg_version, warn
+from ._numba import _numba_thread_limit
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Iterable, KeysView, Mapping
@@ -57,6 +58,7 @@
     "NeighborsView",
     "_choose_graph",
     "_doc_params",
+    "_numba_thread_limit",
     "_resolve_axis",
     "annotate_doc_types",
     "axis_mul_or_truediv",

src/scanpy/_utils/_numba.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import TYPE_CHECKING
+
+import numba
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
+
+@contextmanager
+def _numba_thread_limit(n_threads: int | None) -> Generator[None]:
+    """Temporarily set Numba's thread count and restore it on exit."""
+    if n_threads is None:
+        yield
+        return
+
+    previous = numba.get_num_threads()
+    n_threads = max(1, min(n_threads, numba.config.NUMBA_NUM_THREADS))
+    numba.set_num_threads(n_threads)
+    try:
+        yield
+    finally:
+        numba.set_num_threads(previous)

src/scanpy/tools/_rank_genes_groups.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-from contextlib import contextmanager
 from typing import TYPE_CHECKING
 
 import numba
@@ -18,6 +17,7 @@
 from .._settings import Default
 from .._settings.presets import DETest
 from .._utils import (
+    _numba_thread_limit,
     check_nonnegative_integers,
     get_literal_vals,
     raise_not_implemented_error_if_backed_type,
@@ -47,22 +47,6 @@ def _select_top_n(scores: NDArray, n_top: int):
     return global_indices
 
 
-@contextmanager
-def _numba_thread_limit(n_threads: int | None) -> Generator[None, None, None]:
-    """Temporarily set Numba's thread count and restore it on exit."""
-    if n_threads is None:
-        yield
-        return
-
-    previous = numba.get_num_threads()
-    n_threads = max(1, min(n_threads, numba.config.NUMBA_NUM_THREADS))
-    numba.set_num_threads(n_threads)
-    try:
-        yield
-    finally:
-        numba.set_num_threads(previous)
-
-
 @njit
 def rankdata(data: NDArray[np.number]) -> NDArray[np.float64]:
     """Parallelized version of scipy.stats.rankdata."""

tests/test_rank_genes_groups.py

@@ -4,6 +4,7 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, TypedDict, cast
 
+import numba
 import numpy as np
 import pandas as pd
 import pytest
@@ -16,7 +17,7 @@
 from scanpy._utils import select_groups
 from scanpy.get import rank_genes_groups_df
 from scanpy.tools import rank_genes_groups
-from scanpy.tools._rank_genes_groups import _numba_thread_limit, _RankGenes
+from scanpy.tools._rank_genes_groups import _RankGenes
 from testing.scanpy._helpers import random_mask
 from testing.scanpy._helpers.data import pbmc68k_reduced
 from testing.scanpy._pytest.params import ARRAY_TYPES, ARRAY_TYPES_MEM
@@ -254,82 +255,39 @@ def test_wilcoxon_tie_correction(*, reference: bool) -> None:
     np.testing.assert_allclose(test_obj.stats[groups[0]]["pvals"], pvals, atol=1e-5)
 
 
-def test_wilcoxon_huge_data(monkeypatch):
+def test_wilcoxon_huge_data(monkeypatch: pytest.MonkeyPatch) -> None:
     max_size = 300
     adata = pbmc68k_reduced()
     monkeypatch.setattr(sc.tl._rank_genes_groups, "_CONST_MAX_SIZE", max_size)
     rank_genes_groups(adata, groupby="bulk_labels", method="wilcoxon")
 
 
-def test_numba_thread_limit_restores_previous_value(monkeypatch):
-    calls = []
-    monkeypatch.setattr(sc.tl._rank_genes_groups.numba, "get_num_threads", lambda: 8)
-    monkeypatch.setattr(sc.tl._rank_genes_groups.numba, "set_num_threads", calls.append)
-
-    with _numba_thread_limit(2):
-        pass
-
-    assert calls == [2, 8]
-
-
-def test_numba_thread_limit_restores_previous_value_on_exception(monkeypatch):
-    calls = []
-    msg = "synthetic failure"
-    monkeypatch.setattr(sc.tl._rank_genes_groups.numba, "get_num_threads", lambda: 8)
-    monkeypatch.setattr(sc.tl._rank_genes_groups.numba, "set_num_threads", calls.append)
-
-    with (
-        pytest.raises(RuntimeError, match=msg),
-        _numba_thread_limit(2),
-    ):
-        raise RuntimeError(msg)
-
-    assert calls == [2, 8]
-
-
-def test_numba_thread_limit_clamps_to_configured_maximum(monkeypatch):
-    calls = []
-    monkeypatch.setattr(sc.tl._rank_genes_groups.numba, "get_num_threads", lambda: 3)
-    monkeypatch.setattr(sc.tl._rank_genes_groups.numba, "set_num_threads", calls.append)
-    monkeypatch.setattr(sc.tl._rank_genes_groups.numba.config, "NUMBA_NUM_THREADS", 4)
-
-    with _numba_thread_limit(99):
-        pass
-
-    assert calls[0] == 4
-    assert calls[-1] == 3
-
-
-def test_wilcoxon_sets_numba_threads_from_settings(monkeypatch):
-    calls = []
-    old_n_jobs = sc.settings.n_jobs
-    monkeypatch.setattr(sc.tl._rank_genes_groups.numba, "get_num_threads", lambda: 8)
-    monkeypatch.setattr(sc.tl._rank_genes_groups.numba, "set_num_threads", calls.append)
-
-    try:
-        sc.settings.n_jobs = 2
-        adata = get_example_data(np.asarray)
-        rank_genes_groups(adata, "true_groups", n_genes=5, method="wilcoxon")
-    finally:
-        sc.settings.n_jobs = old_n_jobs
-
-    assert 2 in calls, "Wilcoxon path did not use scanpy.settings.n_jobs."
-    assert calls[-1] == 8
-
-
-def test_t_test_does_not_set_numba_threads_from_settings(monkeypatch):
-    calls = []
+@pytest.mark.parametrize(
+    "method",
+    [
+        pytest.param(
+            "t-test", marks=pytest.mark.xfail(reason="t-test doesn’t use numba (yet)")
+        ),
+        "wilcoxon",
+    ],
+)
+def test_set_numba_threads_from_settings(
+    monkeypatch: pytest.MonkeyPatch, method: Literal["t-test", "wilcoxon"]
+) -> None:
+    was_set_to = []
     old_n_jobs = sc.settings.n_jobs
-    monkeypatch.setattr(sc.tl._rank_genes_groups.numba, "set_num_threads", calls.append)
+    monkeypatch.setattr(numba, "get_num_threads", lambda: 8)
+    monkeypatch.setattr(numba, "set_num_threads", was_set_to.append)
 
     try:
         sc.settings.n_jobs = 2
         adata = get_example_data(np.asarray)
-        rank_genes_groups(adata, "true_groups", n_genes=5, method="t-test")
+        rank_genes_groups(adata, "true_groups", n_genes=5, method=method)
     finally:
         sc.settings.n_jobs = old_n_jobs
 
-    assert calls == []
+    assert 2 in was_set_to, "Wilcoxon path did not use scanpy.settings.n_jobs."
+    assert was_set_to[-1] == 8
 
 
 @pytest.mark.parametrize(

tests/test_utils.py

@@ -2,17 +2,20 @@
 
 import itertools
 import string
+from contextlib import suppress
 from operator import mul, truediv
 from types import ModuleType
 from typing import TYPE_CHECKING
 
+import numba
 import numpy as np
 import pytest
 from anndata.tests.helpers import asarray
 from scipy import sparse
 
 from scanpy._compat import CSBase, DaskArray
 from scanpy._utils import (
+    _numba_thread_limit,
     axis_mul_or_truediv,
     check_nonnegative_integers,
     descend_classes_and_funcs,
@@ -240,3 +243,32 @@ def test_random_str() -> None:
     assert strings.dtype == np.dtype("U2")
     unique = np.unique(strings, axis=0)
     assert len(unique) == len(strings)
+
+
+@pytest.mark.parametrize("success", [True, False], ids=["success", "exception"])
+def test_numba_thread_limit_restores_previous_value(
+    *, monkeypatch: pytest.MonkeyPatch, success: bool
+) -> None:
+    was_set_to = []
+    monkeypatch.setattr(numba, "get_num_threads", lambda: 8)
+    monkeypatch.setattr(numba, "set_num_threads", was_set_to.append)
+
+    with suppress(RuntimeError), _numba_thread_limit(2):
+        if not success:
+            raise RuntimeError
+
+    assert was_set_to == [2, 8]
+
+
+def test_numba_thread_limit_clamps_to_configured_maximum(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    was_set_to = []
+    monkeypatch.setattr(numba, "get_num_threads", lambda: 3)
+    monkeypatch.setattr(numba, "set_num_threads", was_set_to.append)
+    monkeypatch.setattr(numba.config, "NUMBA_NUM_THREADS", 4)
+
+    with _numba_thread_limit(99):
+        pass
+
+    assert was_set_to == [4, 3]