Lazy neighborhood filter kernel compilation

benchmarks/mpas_ocean.py

@@ -107,7 +107,8 @@ def track_nbytes_gradient(self, resolution):
 
     def track_peakmem_gradient(self, resolution):
         """Transient high-water allocation of taking a gradient."""
-        return peak_allocated(lambda: self.uxds[data_var].gradient())
+        with numba_threads(1):
+            return peak_allocated(lambda: self.uxds[data_var].gradient())
 
     track_peakmem_gradient.unit = "bytes"
 

test/test_dependencies.py

@@ -38,4 +38,34 @@ def test_hvplot_optional():
     _assert_not_imported_after_import_uxarray("hvplot")
 
 
+def test_no_numba_kernels_built_on_import():
+    """Test that `import uxarray` does not build any numba kernel.
+
+    ``guvectorize`` compiles at decoration time when it is given explicit
+    signatures, so a kernel assigned at module scope is built during the
+    import. This compilation can dominate the uxarray import, and building a
+    ``target="parallel"`` kernel starts numba's threading layer, which
+    leaves a thread pool running, making forks unsafe.
+    """
+    code = (
+        "import numba, uxarray\n"
+        "try:\n"
+        "    layer = numba.threading_layer()\n"
+        "except ValueError:\n"
+        "    pass\n"
+        "else:\n"
+        "    raise AssertionError(\n"
+        "        f'`import uxarray` started numba threading layer {layer!r}. '\n"
+        "        'Something it imports builds a parallel kernel at module '\n"
+        "        'scope; build it on first use instead.'\n"
+        "    )\n"
+    )
+    result = subprocess.run(
+        [sys.executable, "-c", code],
+        capture_output=True,
+        text=True,
+    )
+    assert result.returncode == 0, result.stderr
+
+
 # TODO: similar tests for cartopy, holoviews, and other optional deps.

uxarray/grid/neighbors.py

@@ -1,3 +1,4 @@
+import functools
 import warnings
 from typing import Callable
 
@@ -1277,26 +1278,6 @@ def _median(window, _):
     return np.median(window)
 
 
-# One compiled kernel per reduction. The methods on ``Neighborhood`` below name
-# these directly, so there is no dispatch table between the public API and the
-# gufuncs: a reduction is reachable only if a method exists for it, and a method
-# can only reach the kernel it names. ``Neighborhood`` is the only class that
-# names them -- the data-bound classes reach a kernel by naming the
-# ``Neighborhood`` method for it, so there is one place per reduction where its
-# kernel and parameter are chosen.
-_MEAN_KERNEL = _make_kernel(lambda window, _: np.mean(window))
-_SUM_KERNEL = _make_kernel(lambda window, _: np.sum(window))
-_MIN_KERNEL = _make_kernel(lambda window, _: np.min(window))
-_MAX_KERNEL = _make_kernel(lambda window, _: np.max(window))
-_PTP_KERNEL = _make_kernel(lambda window, _: np.max(window) - np.min(window))
-_MEDIAN_KERNEL = _make_kernel(_median)
-_VAR_KERNEL = _make_kernel(_variance)
-_STD_KERNEL = _make_kernel(lambda window, ddof: np.sqrt(_variance(window, ddof)))
-# ``percentile`` is ``quantile`` on a 0-100 scale, so both methods rescale onto
-# this one kernel rather than compiling a near-duplicate.
-_QUANTILE_KERNEL = _make_kernel(lambda window, q: np.quantile(window, q))
-
-
 def _as_quantile(q, scale: float):
     """Validates ``q`` on a 0-``scale`` scale and returns it as a 0-1 fraction."""
     value = float(q)
@@ -1507,47 +1488,107 @@ def __repr__(self) -> str:
             f"neighbors_per_element=[{self._counts.min()}, {self._counts.max()}]>"
         )
 
+    # One compiled kernel per reduction. The methods below call into these
+    # directly. Non-compiled functions are only provided hooks through
+    # ``reduce``. If new compiled reductions are desired, they should follow
+    # this pattern.
+    #
+    # ``functools.cache`` defers each build to the 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 ``staticmethod``s rather than attributes for the same
+    # reason: a class body runs at import, so assigning them there would
+    # compile all nine during ``import uxarray``.
+
+    @staticmethod
+    @functools.cache
+    def _mean_kernel():
+        return _make_kernel(lambda window, _: np.mean(window))
+
+    @staticmethod
+    @functools.cache
+    def _sum_kernel():
+        return _make_kernel(lambda window, _: np.sum(window))
+
+    @staticmethod
+    @functools.cache
+    def _min_kernel():
+        return _make_kernel(lambda window, _: np.min(window))
+
+    @staticmethod
+    @functools.cache
+    def _max_kernel():
+        return _make_kernel(lambda window, _: np.max(window))
+
+    @staticmethod
+    @functools.cache
+    def _ptp_kernel():
+        return _make_kernel(lambda window, _: np.max(window) - np.min(window))
+
+    @staticmethod
+    @functools.cache
+    def _median_kernel():
+        return _make_kernel(_median)
+
+    @staticmethod
+    @functools.cache
+    def _var_kernel():
+        return _make_kernel(_variance)
+
+    @staticmethod
+    @functools.cache
+    def _std_kernel():
+        return _make_kernel(lambda window, ddof: np.sqrt(_variance(window, ddof)))
+
+    # ``percentile`` is ``quantile`` on a 0-100 scale, so both methods
+    # rescale onto this one kernel rather than compiling a near-duplicate.
+    @staticmethod
+    @functools.cache
+    def _quantile_kernel():
+        return _make_kernel(lambda window, q: np.quantile(window, q))
+
     def mean(self, uxda):
         """Mean of each neighborhood."""
-        return self._apply_kernel(uxda, _MEAN_KERNEL, 0.0)
+        return self._apply_kernel(uxda, self._mean_kernel, 0.0)
 
     def sum(self, uxda):
         """Sum of each neighborhood."""
-        return self._apply_kernel(uxda, _SUM_KERNEL, 0.0)
+        return self._apply_kernel(uxda, self._sum_kernel, 0.0)
 
     def min(self, uxda):
         """Smallest value in each neighborhood."""
-        return self._apply_kernel(uxda, _MIN_KERNEL, 0.0)
+        return self._apply_kernel(uxda, self._min_kernel, 0.0)
 
     def max(self, uxda):
         """Largest value in each neighborhood."""
-        return self._apply_kernel(uxda, _MAX_KERNEL, 0.0)
+        return self._apply_kernel(uxda, self._max_kernel, 0.0)
 
     def ptp(self, uxda):
         """Peak-to-peak spread (``max - min``) of each neighborhood."""
-        return self._apply_kernel(uxda, _PTP_KERNEL, 0.0)
+        return self._apply_kernel(uxda, self._ptp_kernel, 0.0)
 
     def median(self, uxda):
         """Median of each neighborhood."""
-        return self._apply_kernel(uxda, _MEDIAN_KERNEL, 0.0)
+        return self._apply_kernel(uxda, self._median_kernel, 0.0)
 
     def var(self, uxda, ddof: int = 0):
         """Variance of each neighborhood, with ``ddof`` delta degrees of
         freedom."""
-        return self._apply_kernel(uxda, _VAR_KERNEL, float(ddof))
+        return self._apply_kernel(uxda, self._var_kernel, float(ddof))
 
     def std(self, uxda, ddof: int = 0):
         """Standard deviation of each neighborhood, with ``ddof`` delta degrees
         of freedom."""
-        return self._apply_kernel(uxda, _STD_KERNEL, float(ddof))
+        return self._apply_kernel(uxda, self._std_kernel, float(ddof))
 
     def quantile(self, uxda, q: float):
         """Quantile ``q`` (between 0 and 1) of each neighborhood."""
-        return self._apply_kernel(uxda, _QUANTILE_KERNEL, _as_quantile(q, 1.0))
+        return self._apply_kernel(uxda, self._quantile_kernel, _as_quantile(q, 1.0))
 
     def percentile(self, uxda, q: float):
         """Percentile ``q`` (between 0 and 100) of each neighborhood."""
-        return self._apply_kernel(uxda, _QUANTILE_KERNEL, _as_quantile(q, 100.0))
+        return self._apply_kernel(uxda, self._quantile_kernel, _as_quantile(q, 100.0))
 
     def reduce(self, uxda, func: Callable):
         """Reduces each neighborhood with an arbitrary callable.
@@ -1592,7 +1633,7 @@ def run(block, arrays):
             # path does too by writing into a float64 output.
             if block.dtype not in (np.float64, np.float32):
                 block = block.astype(np.float64)
-            return kernel(block, *arrays, param)
+            return kernel()(block, *arrays, param)
 
         return self._apply(uxda, run)