From c03b0bd71fd24dc85c72af9dbb6718c326ba85e3 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Mon, 14 Sep 2026 13:41:20 -0500 Subject: [PATCH 1/3] Make neighborhood filter kernels cacheable across processes The neighborhood gufuncs are built with ``cache=True``, but the cache never hit: every process recompiled all nine kernels and appended a fresh set of entries to the index, which grew without bound (over a thousand files in a working checkout). The cause is that ``_make_kernel`` closed each kernel over its reducer. Numba keys a cached function on a hash of its closure, and a ``Dispatcher`` serializes with a ``uuid4`` regenerated in every process, so the key differed in every process. Numba reports no error for this -- it writes the cache and silently misses it. All nine kernels also shared one cache file, since the closure made them one function as far as numba's locator was concerned. Give each reduction a module-level kernel body instead, and apply ``guvectorize`` to it on first call. A module-level body has no closure, so its key is stable, and it reaches its reducer as a global, which numba resolves at compile time and leaves out of the key. Each body also gets its own cache file. Kernels are now compiled once per machine rather than once per process. Compilation stays lazy. ``guvectorize`` compiles at decoration time when given explicit signatures, so decorating at module scope would rebuild every kernel during ``import uxarray`` and start numba's threading layer, leaving a thread pool that makes forks unsafe -- the regression ``test_no_numba_kernels_built_on_import`` guards against. Replace ``functools.cache`` with a double-checked lock while here. It does not hold a lock across the call it memoizes, so under ``dask="parallelized"`` every worker thread that reached a kernel before the first build finished started its own full compilation, serialized behind numba's global compiler lock. A 12-chunk reduction compiled the same kernel 12 times; it now compiles once. The bodies are spelled out rather than generated because the deduplication one would reach for -- a single shared body taking the reducer as an argument -- makes the reducer a dynamic global, which numba refuses to cache at all. The gather is shared through ``_widest``/``_gather`` instead, leaving only the reducer name different between bodies. Co-Authored-By: Claude Opus 5 --- uxarray/grid/neighbors.py | 279 ++++++++++++++++++++++++++++++-------- 1 file changed, 220 insertions(+), 59 deletions(-) diff --git a/uxarray/grid/neighbors.py b/uxarray/grid/neighbors.py index eaf1f9b6a..567298535 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -1,4 +1,4 @@ -import functools +import threading import warnings from typing import Callable @@ -1214,49 +1214,104 @@ 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. +def _lazy(impl): + """Returns a kernel that compiles ``impl`` into a gufunc on its first call. + + Two properties of ``impl`` are load-bearing, and both are easy to undo by + accident. + + *It is compiled lazily.* ``guvectorize`` compiles at decoration time when it + is given explicit signatures, so decorating at module scope would build + every kernel during ``import uxarray``. That dominated the import, and for + ``target="parallel"`` it also started numba's threading layer, leaving a + thread pool that makes forks unsafe. + ``test_no_numba_kernels_built_on_import`` guards against a regression. + + *It is a module-level function.* ``cache=True`` needs a real source file to + key against, and, less obviously, the body must not close over anything. + Numba keys a cached function on a hash of its closure, and a ``Dispatcher`` + serializes with a ``uuid4`` regenerated in every process; a kernel closing + over a reducer therefore hashed differently in every process, so the cache + never hit and its index grew without bound. A module-level body has no + closure and reaches its reducer as a global, which numba resolves at + compile time and leaves out of the key, so each kernel is compiled once per + machine instead of once per process. Edits still invalidate the cache: + numba stamps it with the source file's mtime, and the reducers live in this + file. + + The lock is not optional. ``_apply`` hands these kernels to + ``dask="parallelized"``, so the first call can arrive on every worker thread + at once. ``functools.cache`` does not hold a lock across the call it + memoizes, so each thread would start its own full compilation, serialized + behind numba's global compiler lock. """ - # 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) - - return kernel + lock = threading.Lock() + built = [] def kernel(*args): - return build()(*args) + if not built: + with lock: + if not built: + built.append( + guvectorize( + _GUFUNC_SIGNATURES, _GUFUNC_LAYOUT, **_GUFUNC_KWARGS + )(impl) + ) + return built[0](*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 that the kernel bodies +# below can reach it as a global rather than closing over it -- see ``_lazy``. +@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 +1326,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 +1345,114 @@ def _median(window, _): return np.median(window) +# One kernel body per reduction. They are spelled out rather than generated +# because each must be a module-level function with no closure for ``cache=True`` +# to work (see ``_lazy``). The obvious deduplication -- a single shared body +# taking the reducer as an argument -- makes the reducer a dynamic global, which +# numba refuses to cache at all ("Cannot cache compiled function ... as it uses +# dynamic globals"), so the gather is shared through ``_widest``/``_gather`` +# instead and only the reducer named on the last line differs. + + +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 +1668,25 @@ 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))) - ) + # ``_lazy`` defers each build to the kernel's first call, so none of these + # is compiled by ``import uxarray`` and a reduction that is never used is + # never built. That also keeps the import from starting numba's threading + # layer, which would leave a thread pool behind and make forks unsafe. + # 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(_lazy(_mean_impl)) + _sum_kernel = staticmethod(_lazy(_sum_impl)) + _min_kernel = staticmethod(_lazy(_min_impl)) + _max_kernel = staticmethod(_lazy(_max_impl)) + _ptp_kernel = staticmethod(_lazy(_ptp_impl)) + _median_kernel = staticmethod(_lazy(_median_impl)) + _var_kernel = staticmethod(_lazy(_variance_impl)) + _std_kernel = staticmethod(_lazy(_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 = staticmethod(_lazy(_quantile_impl)) def mean(self, uxda): """Mean of each neighborhood.""" From 1ef918ea63e1904d7e4da048a318f4e70f6edb1f Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Tue, 15 Sep 2026 16:15:44 -0500 Subject: [PATCH 2/3] Cache neighborhood kernels on a descriptor instead of a closure Replaces the ``_lazy`` closure with a ``_LazyKernel`` descriptor that holds the compiled gufunc on itself, so each kernel is still built on first use but the memoization lives at class level rather than in a per-kernel closure. Two things fall out. The ``staticmethod`` wrapping goes away: a descriptor hands back the gufunc itself, so ``self`` is never bound as the kernel's first argument, and the comment explaining that wart goes with it. And compilation now happens where the attribute is resolved -- ``_apply_kernel`` reads it on the calling thread -- so a dask-backed reduction builds its kernel while the graph is being assembled rather than inside a task. The array stays lazy; only the compile moves, and it moves out of the parallel region, which makes the worker-thread race structurally impossible rather than merely locked against. ``functools.cached_property`` would not do. It caches per instance, and ``Grid.neighborhood()`` returns a new ``Neighborhood`` on every call, so it ran ``guvectorize`` once per neighborhood: fifty fresh instances measured fifty compilations, against zero for the descriptor. It also holds no lock, CPython having removed it in 3.12, so twelve threads racing one instance measured twelve compilations, against one. The kernel bodies must still be module-level functions with no closure. That is what keeps numba's cache key stable across processes, and it is independent of where the memoization lives -- the memoizer never enters the key. Co-Authored-By: Claude Opus 5 --- uxarray/grid/neighbors.py | 122 +++++++++++++++++--------------------- 1 file changed, 53 insertions(+), 69 deletions(-) diff --git a/uxarray/grid/neighbors.py b/uxarray/grid/neighbors.py index 567298535..4239dfcf0 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -1214,52 +1214,40 @@ def _get_element_coords(grid, data_mapping: str, coordinate_system: str): _GUFUNC_KWARGS = {"nopython": True, "cache": True, "target": "parallel"} -def _lazy(impl): - """Returns a kernel that compiles ``impl`` into a gufunc on its first call. - - Two properties of ``impl`` are load-bearing, and both are easy to undo by - accident. - - *It is compiled lazily.* ``guvectorize`` compiles at decoration time when it - is given explicit signatures, so decorating at module scope would build - every kernel during ``import uxarray``. That dominated the import, and for - ``target="parallel"`` it also started numba's threading layer, leaving a - thread pool that makes forks unsafe. - ``test_no_numba_kernels_built_on_import`` guards against a regression. - - *It is a module-level function.* ``cache=True`` needs a real source file to - key against, and, less obviously, the body must not close over anything. - Numba keys a cached function on a hash of its closure, and a ``Dispatcher`` - serializes with a ``uuid4`` regenerated in every process; a kernel closing - over a reducer therefore hashed differently in every process, so the cache - never hit and its index grew without bound. A module-level body has no - closure and reaches its reducer as a global, which numba resolves at - compile time and leaves out of the key, so each kernel is compiled once per - machine instead of once per process. Edits still invalidate the cache: - numba stamps it with the source file's mtime, and the reducers live in this - file. - - The lock is not optional. ``_apply`` hands these kernels to - ``dask="parallelized"``, so the first call can arrive on every worker thread - at once. ``functools.cache`` does not hold a lock across the call it - memoizes, so each thread would start its own full compilation, serialized - behind numba's global compiler lock. +class _LazyKernel: + """Compiles ``impl`` into a gufunc on first access, once per process. + + ``impl`` must stay a module-level function, for two reasons that fail + quietly. ``guvectorize`` compiles at decoration time, so building one at + module scope would compile during ``import uxarray`` and start numba's + threading layer, leaving a thread pool that makes forks unsafe + (``test_no_numba_kernels_built_on_import`` guards this); and 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 and ``cache=True`` never hits. 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 one. """ - lock = threading.Lock() - built = [] - - def kernel(*args): - if not built: - with lock: - if not built: - built.append( - guvectorize( - _GUFUNC_SIGNATURES, _GUFUNC_LAYOUT, **_GUFUNC_KWARGS - )(impl) - ) - return built[0](*args) - return kernel + 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 @njit(cache=True) @@ -1280,8 +1268,8 @@ def _gather(data, flat, start, count, buffer): # Reducers take ``(window, param)``; those without a parameter ignore the -# second argument. Each is a module-level ``njit`` so that the kernel bodies -# below can reach it as a global rather than closing over it -- see ``_lazy``. +# 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) @@ -1345,13 +1333,11 @@ def _median(window, _): return np.median(window) -# One kernel body per reduction. They are spelled out rather than generated -# because each must be a module-level function with no closure for ``cache=True`` -# to work (see ``_lazy``). The obvious deduplication -- a single shared body -# taking the reducer as an argument -- makes the reducer a dynamic global, which -# numba refuses to cache at all ("Cannot cache compiled function ... as it uses -# dynamic globals"), so the gather is shared through ``_widest``/``_gather`` -# instead and only the reducer named on the last line differs. +# 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): @@ -1668,25 +1654,23 @@ def __repr__(self) -> str: # ``reduce``. If new compiled reductions are desired, they should follow # this pattern. # - # ``_lazy`` defers each build to the kernel's first call, so none of these - # is compiled by ``import uxarray`` and a reduction that is never used is - # never built. That also keeps the import from starting numba's threading - # layer, which would leave a thread pool behind and make forks unsafe. - # 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(_lazy(_mean_impl)) - _sum_kernel = staticmethod(_lazy(_sum_impl)) - _min_kernel = staticmethod(_lazy(_min_impl)) - _max_kernel = staticmethod(_lazy(_max_impl)) - _ptp_kernel = staticmethod(_lazy(_ptp_impl)) - _median_kernel = staticmethod(_lazy(_median_impl)) - _var_kernel = staticmethod(_lazy(_variance_impl)) - _std_kernel = staticmethod(_lazy(_std_impl)) + # 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(_lazy(_quantile_impl)) + _quantile_kernel = _LazyKernel(_quantile_impl) def mean(self, uxda): """Mean of each neighborhood.""" From 4d7db90c501412d2d35880c263946efdaf14c4e2 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Wed, 16 Sep 2026 17:28:32 -0500 Subject: [PATCH 3/3] nb kernels comment cleanup --- uxarray/grid/neighbors.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/uxarray/grid/neighbors.py b/uxarray/grid/neighbors.py index 4239dfcf0..8df9e49c9 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -1217,20 +1217,19 @@ def _get_element_coords(grid, data_mapping: str, coordinate_system: str): class _LazyKernel: """Compiles ``impl`` into a gufunc on first access, once per process. - ``impl`` must stay a module-level function, for two reasons that fail - quietly. ``guvectorize`` compiles at decoration time, so building one at - module scope would compile during ``import uxarray`` and start numba's - threading layer, leaving a thread pool that makes forks unsafe - (``test_no_numba_kernels_built_on_import`` guards this); and numba keys its + ``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 and ``cache=True`` never hits. Holding the gufunc on the - descriptor rather than the instance keeps the throwaway ``Neighborhood`` + 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 one. + for inside of one. """ def __init__(self, impl):