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/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..753092f18 100644 --- a/uxarray/grid/neighbors.py +++ b/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)