fix: numba cache data files named by key, so concurrent compiles can't mix up signatures

src/infercnvpy/_numba.py

@@ -3,29 +3,69 @@
 A ``parallel=True`` kernel only runs on the main thread. Called from any other thread (a dask worker, a thread pool),
 the serial twin runs instead: numba's default ``workqueue`` threading layer (the only one in the macOS wheels) aborts
 the process on concurrent use, and a caller that brings its own threads should not get numba threads on top.
+
+All kernels use numba's disk cache, with data files named after their cache key (see `_KeyedCacheFile`).
 """
 
 from __future__ import annotations
 
+import hashlib
 import threading
 from contextlib import contextmanager
 from functools import update_wrapper
 from types import FunctionType
 from typing import TYPE_CHECKING, Any
 
 import numba
+from numba.core.caching import FunctionCache, IndexDataCacheFile
 
 if TYPE_CHECKING:
     from builtins import range as prange  # numba's prange iterates like range; its own type says nothing about that
     from collections.abc import Callable, Generator
 else:
     from numba import prange
 
-__all__ = ["n_kernel_chunks", "njit_parallel", "numba_threads", "prange"]
+__all__ = ["n_kernel_chunks", "njit", "njit_parallel", "numba_threads", "prange"]
+
+
+class _KeyedCacheFile(IndexDataCacheFile):
+    """numba's cache index and data files, with each data file named after its key (signature, CPU, bytecode).
+
+    numba numbers the data files of a function (``.1.nbc``, ``.2.nbc``, ...) by the free slots in the index as the
+    process read it. Processes that compile different signatures of one function at the same time (pytest-xdist
+    workers, a pool of jobs on one environment) read the same index and pick the same number. The index can then end
+    up mapping one signature to another signature's machine code, and every later call with it fails with "can't
+    unbox array from PyObject" until the cache is deleted. With a name per key, a data file only ever holds the code
+    for its own key. Concurrent index writes can still drop an entry, which only costs one recompilation.
+    """
+
+    def save(self, key, data):
+        digest = hashlib.sha256(repr(key).encode()).hexdigest()[:16]
+        data_name = self._data_name_pattern.replace("{number:d}", digest)
+        self._save_data(data_name, data)  # before the index, so that no index entry points to a missing file
+        overloads = self._load_index()
+        if overloads.get(key) != data_name:
+            overloads[key] = data_name
+            self._save_index(overloads)
+
+
+class _KeyedFunctionCache(FunctionCache):
+    def __init__(self, py_func):
+        super().__init__(py_func)
+        self._cache_file = _KeyedCacheFile(
+            self._cache_path, self._impl.filename_base, self._impl.locator.get_source_stamp()
+        )
+
+
+def njit(fn: Callable[..., Any], /, **options: Any) -> Any:
+    """`numba.njit` with ``nogil=True`` and the disk cache of `_KeyedCacheFile`."""
+    dispatcher = numba.njit(fn, nogil=True, **options)
+    dispatcher._cache = _KeyedFunctionCache(dispatcher.py_func)  # type: ignore[attr-defined]  # what `cache=True` sets
+    return dispatcher
 
 
 class njit_parallel:
-    """Compile a function with numba twice, ``parallel=True`` and serial, both ``cache=True`` and ``nogil=True``.
+    """Compile a function with `njit` twice, ``parallel=True`` and serial.
 
     Calling it runs the parallel version on the main thread and the serial one on any other thread, or always the
     serial one with ``serial=True`` (for code that runs inside dask tasks). The two versions are distinct function
@@ -34,8 +74,8 @@ class njit_parallel:
 
     def __init__(self, fn: Callable[..., Any], /):
         assert isinstance(fn, FunctionType)
-        self.parallel = numba.njit(_renamed(fn, "parallel"), parallel=True, cache=True, nogil=True)
-        self.serial = numba.njit(_renamed(fn, "serial"), cache=True, nogil=True)
+        self.parallel = njit(_renamed(fn, "parallel"), parallel=True)
+        self.serial = njit(_renamed(fn, "serial"))
         update_wrapper(self, fn)
 
     def __call__(self, *args, serial: bool = False):

src/infercnvpy/tl/_infercnv_kernels.py

@@ -20,13 +20,12 @@
 
 from __future__ import annotations
 
-import numba
 import numpy as np
 
-from infercnvpy._numba import njit_parallel, prange
+from infercnvpy._numba import njit, njit_parallel, prange
 
 
-@numba.njit(cache=True, nogil=True)
+@njit
 def _lfc(v, lo, hi, cap):
     """Clipped log fold change of `v` against the reference bounds `[lo, hi]` (`lo == hi` for one reference)."""
     if v > hi:
@@ -42,7 +41,7 @@ def _lfc(v, lo, hi, cap):
     return d
 
 
-@numba.njit(cache=True, nogil=True)
+@njit
 def _select(a, n, k):
     """Partially sort `a[:n]` in place so that `a[k]` is its k-th smallest value, and return that (quickselect)."""
     lo = 0
@@ -76,7 +75,7 @@ def _select(a, n, k):
     return a[k]
 
 
-@numba.njit(cache=True, nogil=True)
+@njit
 def _median(values, buf):
     """Median of `values` (as `np.median`), using `buf` (same size) as scratch."""
     n = values.size
@@ -93,7 +92,7 @@ def _median(values, buf):
     return (lower + upper) / 2
 
 
-@numba.njit(cache=True, nogil=True)
+@njit
 def _smooth(x, p1, p2, win, g0s, gs, As, Bs, Ks, w0s, step):
     """Running means of the kept windows of one cell (gene-ordered `x`) into `win`, from prefix sums."""
     for c in range(g0s.size):
@@ -114,7 +113,7 @@ def _smooth(x, p1, p2, win, g0s, gs, As, Bs, Ks, w0s, step):
             win[w0 + w] = (p2[s + b] - p2[s]) / weight
 
 
-@numba.njit(cache=True, nogil=True)
+@njit
 def _scatter(win, j, d, first, weights):
     """Add `d` times the pyramid weights of gene `j` to the weighted window sums `win`, from window `first[j]` on.
 
@@ -126,14 +125,14 @@ def _scatter(win, j, d, first, weights):
         win[w + m] += d * wj[m]
 
 
-@numba.njit(cache=True, nogil=True)
+@njit
 def _normalise(win, n_windows, window_weight):
     """Divide the weighted window sums by the sums of their weights: the running means."""
     for w in range(n_windows):
         win[w] /= window_weight[w]
 
 
-@numba.njit(cache=True, nogil=True)
+@njit
 def _center_emit(win, buf, thr, centre, out, r, rowstats, counts):
     """Subtract the median from the running means `win`, write them to `out[r]` as float32.
 

src/infercnvpy/tl/_spatial.py

@@ -13,18 +13,17 @@
 
 from typing import Any
 
-import numba
 import numpy as np
 import pandas as pd
 import scipy.sparse as sp
 from anndata import AnnData
 
 from infercnvpy._logging import log, log_result
-from infercnvpy._numba import njit_parallel, numba_threads, prange
+from infercnvpy._numba import njit, njit_parallel, numba_threads, prange
 from infercnvpy._util import _as_matrix, _params_for_uns
 
 
-@numba.njit(cache=True, nogil=True)
+@njit
 def _row_weights(g_indptr, g_indices, g_data, i, alpha):
     """Self weight `a` and neighbour scale `(1 - a) / d` of cell `i` (diagonal entries of the graph are ignored)."""
     d = 0.0
@@ -73,7 +72,7 @@ def _hop_dense(g_indptr, g_indices, g_data, alpha, X, out):  # pragma: no cover
                     row[c] += w * X[j, c]
 
 
-@numba.njit(cache=True, nogil=True)
+@njit
 def _icm(indptr, indices, weights, unary, labels, beta, n_iter):  # pragma: no cover
     """Iterated conditional modes for a Potts model, in place on `labels` (-1 = unlabelled). Returns the sweeps run.
 

tests/test_numba_cache.py

@@ -0,0 +1,73 @@
+"""The numba disk cache stays correct when processes compile different signatures of one kernel at the same time."""
+
+import os
+import subprocess
+import sys
+import textwrap
+
+KERNEL = """
+from infercnvpy._numba import njit_parallel
+
+@njit_parallel
+def total(x):
+    s = 0.0
+    for v in x:
+        s += v
+    return s
+"""
+
+# Two processes A and B compile `total` for float32 and float64 at the same time, in this order: A writes the index,
+# B (which read the index before A wrote it) writes the index and its data, then A's data write lands.
+# One process plays both, with two copies of the kernel module and numba's file access patched to that order.
+RACE = """
+import importlib.util
+import sys
+
+import numpy as np
+from numba.core.caching import IndexDataCacheFile
+
+
+def kernel():
+    spec = importlib.util.spec_from_file_location("kernels", sys.argv[1])
+    module = sys.modules["kernels"] = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module.total
+
+
+a, b = kernel(), kernel()
+save_data, load_index = IndexDataCacheFile._save_data, IndexDataCacheFile._load_index
+late = []
+IndexDataCacheFile._save_data = lambda self, name, data: late.append((self, name, data))
+a(np.ones(3, dtype=np.float32))
+IndexDataCacheFile._save_data = save_data
+IndexDataCacheFile._load_index = lambda self: {}
+b(np.ones(3))
+IndexDataCacheFile._load_index = load_index
+for args in late:
+    save_data(*args)
+"""
+
+CALL = """
+import numpy as np
+
+import kernels as module
+
+assert module.total(np.ones(3, dtype=np.float32)) == 3
+assert module.total(np.ones(3)) == 3
+"""
+
+
+def test_concurrent_compiles_share_the_cache(tmp_path):
+    kernels = tmp_path / "kernels.py"
+    kernels.write_text(KERNEL)
+    env = {**os.environ, "NUMBA_CACHE_DIR": str(tmp_path / "cache")}
+    for script in (RACE, CALL):
+        res = subprocess.run(
+            [sys.executable, "-c", textwrap.dedent(script), str(kernels)],
+            env=env,
+            cwd=tmp_path,
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+        assert res.returncode == 0, res.stderr