From ec83767c568e9c96cda504b75ca9748c5da26599 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Tue, 25 Aug 2026 15:30:30 -0500 Subject: [PATCH 1/3] Lazy neighborhood filter kernel compilation --- test/test_dependencies.py | 30 +++++++++++++++++++ uxarray/grid/neighbors.py | 62 +++++++++++++++++++++++++++++++++------ 2 files changed, 83 insertions(+), 9 deletions(-) diff --git a/test/test_dependencies.py b/test/test_dependencies.py index fb27bfbce..2be243d01 100644 --- a/test/test_dependencies.py +++ b/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. diff --git a/uxarray/grid/neighbors.py b/uxarray/grid/neighbors.py index 03a4c7cd5..a0c301b18 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -1,3 +1,4 @@ +import threading import warnings from typing import Callable @@ -1246,6 +1247,43 @@ def kernel(data, flat, starts, counts, param, out): return kernel +class _LazyKernel: + """A kernel that builds itself the first time it is called. + + ``guvectorize`` compiles at decoration time when it is given explicit + signatures, so calling :func:`_make_kernel` at module scope would compile + every reduction during ``import uxarray`` + """ + + __slots__ = ("_reduce_fn", "_kernel", "_lock") + + def __init__(self, reduce_fn): + self._reduce_fn = reduce_fn + self._kernel = None + self._lock = threading.Lock() + + def __call__(self, *args): + kernel = self._kernel + if kernel is None: + # Checked again under the lock: dask's threaded scheduler can call + # one reduction from several workers at once + with self._lock: + kernel = self._kernel + if kernel is None: + kernel = self._kernel = _make_kernel(self._reduce_fn) + return kernel(*args) + + def __getstate__(self): + # Only the reducer travels. A compiled gufunc is a ``numpy.ufunc``, + # which pickle cannot address by name and so refuses outright. + return self._reduce_fn + + def __setstate__(self, reduce_fn): + self._reduce_fn = reduce_fn + self._kernel = None + self._lock = threading.Lock() + + # 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. @@ -1284,17 +1322,23 @@ def _median(window, _): # 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))) +# +# Naming a kernel is what binds a reduction to it; building it is deferred to +# the first call, for the reasons in :class:`_LazyKernel`. The two are +# separate on purpose -- each name below still resolves to its own object when +# this module is read, so a reduction that names no kernel is a ``NameError`` +# here rather than a lookup that fails once someone runs it. +_MEAN_KERNEL = _LazyKernel(lambda window, _: np.mean(window)) +_SUM_KERNEL = _LazyKernel(lambda window, _: np.sum(window)) +_MIN_KERNEL = _LazyKernel(lambda window, _: np.min(window)) +_MAX_KERNEL = _LazyKernel(lambda window, _: np.max(window)) +_PTP_KERNEL = _LazyKernel(lambda window, _: np.max(window) - np.min(window)) +_MEDIAN_KERNEL = _LazyKernel(_median) +_VAR_KERNEL = _LazyKernel(_variance) +_STD_KERNEL = _LazyKernel(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)) +_QUANTILE_KERNEL = _LazyKernel(lambda window, q: np.quantile(window, q)) def _as_quantile(q, scale: float): From 856441007de9366f08df5d64286302c32e94c136 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Tue, 25 Aug 2026 15:58:47 -0500 Subject: [PATCH 2/3] Lazy nb kernels with functools cache instead --- uxarray/grid/neighbors.py | 135 ++++++++++++++++++-------------------- 1 file changed, 65 insertions(+), 70 deletions(-) diff --git a/uxarray/grid/neighbors.py b/uxarray/grid/neighbors.py index a0c301b18..d796f6d78 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -1,4 +1,4 @@ -import threading +import functools import warnings from typing import Callable @@ -1247,43 +1247,6 @@ def kernel(data, flat, starts, counts, param, out): return kernel -class _LazyKernel: - """A kernel that builds itself the first time it is called. - - ``guvectorize`` compiles at decoration time when it is given explicit - signatures, so calling :func:`_make_kernel` at module scope would compile - every reduction during ``import uxarray`` - """ - - __slots__ = ("_reduce_fn", "_kernel", "_lock") - - def __init__(self, reduce_fn): - self._reduce_fn = reduce_fn - self._kernel = None - self._lock = threading.Lock() - - def __call__(self, *args): - kernel = self._kernel - if kernel is None: - # Checked again under the lock: dask's threaded scheduler can call - # one reduction from several workers at once - with self._lock: - kernel = self._kernel - if kernel is None: - kernel = self._kernel = _make_kernel(self._reduce_fn) - return kernel(*args) - - def __getstate__(self): - # Only the reducer travels. A compiled gufunc is a ``numpy.ufunc``, - # which pickle cannot address by name and so refuses outright. - return self._reduce_fn - - def __setstate__(self, reduce_fn): - self._reduce_fn = reduce_fn - self._kernel = None - self._lock = threading.Lock() - - # 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. @@ -1315,30 +1278,62 @@ 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. +# One compiled kernel per reduction. The methods on ``Neighborhood`` below call +# into these directly. Non-compiled functions are only provided hooks through +# ``Neighborhood.reduce``. If new compiled reductions are desired, they should +# follow this pattern. # -# Naming a kernel is what binds a reduction to it; building it is deferred to -# the first call, for the reasons in :class:`_LazyKernel`. The two are -# separate on purpose -- each name below still resolves to its own object when -# this module is read, so a reduction that names no kernel is a ``NameError`` -# here rather than a lookup that fails once someone runs it. -_MEAN_KERNEL = _LazyKernel(lambda window, _: np.mean(window)) -_SUM_KERNEL = _LazyKernel(lambda window, _: np.sum(window)) -_MIN_KERNEL = _LazyKernel(lambda window, _: np.min(window)) -_MAX_KERNEL = _LazyKernel(lambda window, _: np.max(window)) -_PTP_KERNEL = _LazyKernel(lambda window, _: np.max(window) - np.min(window)) -_MEDIAN_KERNEL = _LazyKernel(_median) -_VAR_KERNEL = _LazyKernel(_variance) -_STD_KERNEL = _LazyKernel(lambda window, ddof: np.sqrt(_variance(window, ddof))) +# ``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. + + +@functools.cache +def _mean_kernel(): + return _make_kernel(lambda window, _: np.mean(window)) + + +@functools.cache +def _sum_kernel(): + return _make_kernel(lambda window, _: np.sum(window)) + + +@functools.cache +def _min_kernel(): + return _make_kernel(lambda window, _: np.min(window)) + + +@functools.cache +def _max_kernel(): + return _make_kernel(lambda window, _: np.max(window)) + + +@functools.cache +def _ptp_kernel(): + return _make_kernel(lambda window, _: np.max(window) - np.min(window)) + + +@functools.cache +def _median_kernel(): + return _make_kernel(_median) + + +@functools.cache +def _var_kernel(): + return _make_kernel(_variance) + + +@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. -_QUANTILE_KERNEL = _LazyKernel(lambda window, q: np.quantile(window, q)) +@functools.cache +def _quantile_kernel(): + return _make_kernel(lambda window, q: np.quantile(window, q)) def _as_quantile(q, scale: float): @@ -1553,45 +1548,45 @@ def __repr__(self) -> str: def mean(self, uxda): """Mean of each neighborhood.""" - return self._apply_kernel(uxda, _MEAN_KERNEL, 0.0) + return self._apply_kernel(uxda, _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, _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, _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, _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, _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, _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, _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, _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, _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, _quantile_kernel, _as_quantile(q, 100.0)) def reduce(self, uxda, func: Callable): """Reduces each neighborhood with an arbitrary callable. @@ -1636,7 +1631,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) From 7127726beb6ec7a36740730b409df1ba23a62fd7 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Wed, 26 Aug 2026 12:34:36 -0500 Subject: [PATCH 3/3] Move kernels inside --- benchmarks/mpas_ocean.py | 3 +- uxarray/grid/neighbors.py | 138 +++++++++++++++++++------------------- 2 files changed, 72 insertions(+), 69 deletions(-) diff --git a/benchmarks/mpas_ocean.py b/benchmarks/mpas_ocean.py index adf679a3a..04325563f 100644 --- a/benchmarks/mpas_ocean.py +++ b/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" diff --git a/uxarray/grid/neighbors.py b/uxarray/grid/neighbors.py index d796f6d78..753092f18 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -1278,64 +1278,6 @@ def _median(window, _): return np.median(window) -# One compiled kernel per reduction. The methods on ``Neighborhood`` below call -# into these directly. Non-compiled functions are only provided hooks through -# ``Neighborhood.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. - - -@functools.cache -def _mean_kernel(): - return _make_kernel(lambda window, _: np.mean(window)) - - -@functools.cache -def _sum_kernel(): - return _make_kernel(lambda window, _: np.sum(window)) - - -@functools.cache -def _min_kernel(): - return _make_kernel(lambda window, _: np.min(window)) - - -@functools.cache -def _max_kernel(): - return _make_kernel(lambda window, _: np.max(window)) - - -@functools.cache -def _ptp_kernel(): - return _make_kernel(lambda window, _: np.max(window) - np.min(window)) - - -@functools.cache -def _median_kernel(): - return _make_kernel(_median) - - -@functools.cache -def _var_kernel(): - return _make_kernel(_variance) - - -@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. -@functools.cache -def _quantile_kernel(): - return _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) @@ -1546,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.