Neighborhood filter kernels caching/race condition fixes

uxarray/grid/neighbors.py

@@ -1,4 +1,4 @@
-import functools
+import threading
 import warnings
 from typing import Callable
 
@@ -1214,49 +1214,91 @@ def _get_element_coords(grid, data_mapping: str, coordinate_system: str):
 _GUFUNC_KWARGS = {"nopython": True, "cache": True, "target": "parallel"}
 
 
-def _make_kernel(reduce_fn):
-    """Returns a kernel, compiled on its first call, that gathers each
-    neighborhood, then calls ``reduce_fn(window, param)`` on the 1-D result.
-
-    ``reduce_fn`` must be numba-compilable, and must be defined in a real
-    source file for ``cache=True`` to find it.
+class _LazyKernel:
+    """Compiles ``impl`` into a gufunc on first access, once per process.
+
+    ``impl`` must stay a module-level function, for two reasons: 1) ``guvectorize``
+    compiles at decoration, so building one at module scope would compile during
+    ``import uxarray``, spawning numba's thread pool and making forks unsafe
+    (``test_no_numba_kernels_built_on_import`` guards this); 2) numba keys its
+    cache on a hash of the closure, where a ``Dispatcher`` serializes with a
+    per-process ``uuid4``, so a body capturing its reducer hashes differently in
+    every process. Holding the gufunc on the descriptor
+     rather than the instance keeps the throwaway ``Neighborhood``
+    that ``Grid.neighborhood()`` returns from recompiling, and the explicit lock
+    is why this is not a ``functools.cached_property``, which holds none --
+    though in practice ``_apply_kernel`` resolves the attribute on the calling
+    thread, so the kernel is built before any dask task runs rather than raced
+    for inside of one.
     """
-    # A reducer shared between kernels arrives already compiled; numba rejects
-    # jitting a dispatcher twice.
-    if not hasattr(reduce_fn, "py_func"):
-        reduce_fn = njit(cache=True)(reduce_fn)
-
-    @functools.cache
-    def build():
-        @guvectorize(_GUFUNC_SIGNATURES, _GUFUNC_LAYOUT, **_GUFUNC_KWARGS)
-        def kernel(data, flat, starts, counts, param, out):
-            widest = 0
-            for i in range(counts.shape[0]):
-                if counts[i] > widest:
-                    widest = counts[i]
-            buffer = np.empty(widest, dtype=np.float64)
-
-            for i in range(starts.shape[0]):
-                count = counts[i]
-                if count == 0:
-                    out[i] = np.nan
-                    continue
-                start = starts[i]
-                for j in range(count):
-                    buffer[j] = data[flat[start + j]]
-                out[i] = reduce_fn(buffer[:count], param)
 
+    def __init__(self, impl):
+        self._impl = impl
+        self._lock = threading.Lock()
+        self._kernel = None
+
+    def __get__(self, obj, objtype=None):
+        kernel = self._kernel
+        if kernel is None:
+            with self._lock:
+                if self._kernel is None:
+                    self._kernel = guvectorize(
+                        _GUFUNC_SIGNATURES, _GUFUNC_LAYOUT, **_GUFUNC_KWARGS
+                    )(self._impl)
+                kernel = self._kernel
         return kernel
 
-    def kernel(*args):
-        return build()(*args)
 
-    return kernel
+@njit(cache=True)
+def _widest(counts):
+    """Largest neighborhood, so each kernel allocates its buffer once."""
+    widest = 0
+    for i in range(counts.shape[0]):
+        if counts[i] > widest:
+            widest = counts[i]
+    return widest
+
+
+@njit(cache=True)
+def _gather(data, flat, start, count, buffer):
+    """Copies one neighborhood's values into ``buffer[:count]``."""
+    for j in range(count):
+        buffer[j] = data[flat[start + j]]
 
 
 # Reducers take ``(window, param)``; those without a parameter ignore the
-# second argument. Numba keys its cache by code object rather than qualified
-# name, so the identically-named lambdas below do not collide.
+# second argument. Each is a module-level ``njit`` so the bodies below reach it
+# as a global rather than closing over it -- see ``_LazyKernel``.
+@njit(cache=True)
+def _mean(window, _):
+    return np.mean(window)
+
+
+@njit(cache=True)
+def _sum(window, _):
+    return np.sum(window)
+
+
+@njit(cache=True)
+def _min(window, _):
+    return np.min(window)
+
+
+@njit(cache=True)
+def _max(window, _):
+    return np.max(window)
+
+
+@njit(cache=True)
+def _ptp(window, _):
+    return np.max(window) - np.min(window)
+
+
+@njit(cache=True)
+def _quantile(window, q):
+    return np.quantile(window, q)
+
+
 @njit(cache=True)
 def _variance(window, ddof):
     """Variance with a delta degrees of freedom. Numba's ``np.var`` takes no
@@ -1271,6 +1313,11 @@ def _variance(window, ddof):
     return total / denominator
 
 
+@njit(cache=True)
+def _std(window, ddof):
+    return np.sqrt(_variance(window, ddof))
+
+
 @njit(cache=True)
 def _median(window, _):
     """numba's ``np.median`` selects by partitioning, and whether a NaN survives
@@ -1285,6 +1332,112 @@ def _median(window, _):
     return np.median(window)
 
 
+# One kernel body per reduction, spelled out because each must be a module-level
+# function with no closure (see ``_LazyKernel``). Deduplicating them into one
+# body taking the reducer as an argument makes it a dynamic global, which numba
+# refuses to cache at all, so only the gather is shared and the bodies differ
+# just in the reducer named on the last line.
+
+
+def _mean_impl(data, flat, starts, counts, param, out):
+    buffer = np.empty(_widest(counts), dtype=np.float64)
+    for i in range(starts.shape[0]):
+        count = counts[i]
+        if count == 0:
+            out[i] = np.nan
+            continue
+        _gather(data, flat, starts[i], count, buffer)
+        out[i] = _mean(buffer[:count], param)
+
+
+def _sum_impl(data, flat, starts, counts, param, out):
+    buffer = np.empty(_widest(counts), dtype=np.float64)
+    for i in range(starts.shape[0]):
+        count = counts[i]
+        if count == 0:
+            out[i] = np.nan
+            continue
+        _gather(data, flat, starts[i], count, buffer)
+        out[i] = _sum(buffer[:count], param)
+
+
+def _min_impl(data, flat, starts, counts, param, out):
+    buffer = np.empty(_widest(counts), dtype=np.float64)
+    for i in range(starts.shape[0]):
+        count = counts[i]
+        if count == 0:
+            out[i] = np.nan
+            continue
+        _gather(data, flat, starts[i], count, buffer)
+        out[i] = _min(buffer[:count], param)
+
+
+def _max_impl(data, flat, starts, counts, param, out):
+    buffer = np.empty(_widest(counts), dtype=np.float64)
+    for i in range(starts.shape[0]):
+        count = counts[i]
+        if count == 0:
+            out[i] = np.nan
+            continue
+        _gather(data, flat, starts[i], count, buffer)
+        out[i] = _max(buffer[:count], param)
+
+
+def _ptp_impl(data, flat, starts, counts, param, out):
+    buffer = np.empty(_widest(counts), dtype=np.float64)
+    for i in range(starts.shape[0]):
+        count = counts[i]
+        if count == 0:
+            out[i] = np.nan
+            continue
+        _gather(data, flat, starts[i], count, buffer)
+        out[i] = _ptp(buffer[:count], param)
+
+
+def _median_impl(data, flat, starts, counts, param, out):
+    buffer = np.empty(_widest(counts), dtype=np.float64)
+    for i in range(starts.shape[0]):
+        count = counts[i]
+        if count == 0:
+            out[i] = np.nan
+            continue
+        _gather(data, flat, starts[i], count, buffer)
+        out[i] = _median(buffer[:count], param)
+
+
+def _variance_impl(data, flat, starts, counts, param, out):
+    buffer = np.empty(_widest(counts), dtype=np.float64)
+    for i in range(starts.shape[0]):
+        count = counts[i]
+        if count == 0:
+            out[i] = np.nan
+            continue
+        _gather(data, flat, starts[i], count, buffer)
+        out[i] = _variance(buffer[:count], param)
+
+
+def _std_impl(data, flat, starts, counts, param, out):
+    buffer = np.empty(_widest(counts), dtype=np.float64)
+    for i in range(starts.shape[0]):
+        count = counts[i]
+        if count == 0:
+            out[i] = np.nan
+            continue
+        _gather(data, flat, starts[i], count, buffer)
+        out[i] = _std(buffer[:count], param)
+
+
+def _quantile_impl(data, flat, starts, counts, param, out):
+    buffer = np.empty(_widest(counts), dtype=np.float64)
+    for i in range(starts.shape[0]):
+        count = counts[i]
+        if count == 0:
+            out[i] = np.nan
+            continue
+        _gather(data, flat, starts[i], count, buffer)
+        out[i] = _quantile(buffer[:count], param)
+
+
 def _as_quantile(q, scale: float):
     """Validates ``q`` on a 0-``scale`` scale and returns it as a 0-1 fraction."""
     value = float(q)
@@ -1500,32 +1653,23 @@ def __repr__(self) -> str:
     # ``reduce``. If new compiled reductions are desired, they should follow
     # this pattern.
     #
-    # ``_make_kernel`` defers each build to the kernel's first call. The
-    # deferred compilation ensures that these kernels will only be compiled
-    # individually and lazily. Further, the lazy compilation prevents gufuncs
-    # from spawning threadpools eagerly and disrupting threading and forking in
-    # other contexts. They are wrapped in ``staticmethod`` because a plain
-    # function in a class body would bind ``self`` as the kernel's first
-    # argument.
-
-    _mean_kernel = staticmethod(_make_kernel(lambda window, _: np.mean(window)))
-    _sum_kernel = staticmethod(_make_kernel(lambda window, _: np.sum(window)))
-    _min_kernel = staticmethod(_make_kernel(lambda window, _: np.min(window)))
-    _max_kernel = staticmethod(_make_kernel(lambda window, _: np.max(window)))
-    _ptp_kernel = staticmethod(
-        _make_kernel(lambda window, _: np.max(window) - np.min(window))
-    )
-    _median_kernel = staticmethod(_make_kernel(_median))
-    _var_kernel = staticmethod(_make_kernel(_variance))
-    _std_kernel = staticmethod(
-        _make_kernel(lambda window, ddof: np.sqrt(_variance(window, ddof)))
-    )
+    # Built on first access, so the import compiles nothing and an unused
+    # reduction is never compiled at all. No ``staticmethod`` is needed: a
+    # descriptor hands back the gufunc itself, so ``self`` is never bound as the
+    # kernel's first argument.
+
+    _mean_kernel = _LazyKernel(_mean_impl)
+    _sum_kernel = _LazyKernel(_sum_impl)
+    _min_kernel = _LazyKernel(_min_impl)
+    _max_kernel = _LazyKernel(_max_impl)
+    _ptp_kernel = _LazyKernel(_ptp_impl)
+    _median_kernel = _LazyKernel(_median_impl)
+    _var_kernel = _LazyKernel(_variance_impl)
+    _std_kernel = _LazyKernel(_std_impl)
 
     # ``percentile`` is ``quantile`` on a 0-100 scale, so both methods
     # rescale onto this one kernel rather than compiling a near-duplicate.
-    _quantile_kernel = staticmethod(
-        _make_kernel(lambda window, q: np.quantile(window, q))
-    )
+    _quantile_kernel = _LazyKernel(_quantile_impl)
 
     def mean(self, uxda):
         """Mean of each neighborhood."""