Neighborhood filter kernels caching/race condition fixes - #1768
Merged
Merged
Conversation
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
cmdupuis3
added
bug
ASV BenchmarkingBenchmark Comparison ResultsBenchmarks that have improved:
Benchmarks that have stayed the same:
Benchmarks that have got worse:
|
cmdupuis3
marked this pull request as ready for review
September 18, 2026 21:50
Sevans711
reviewed
Sep 23, 2026
Sevans711
left a comment
Sevans711
left a comment
Collaborator
There was a problem hiding this comment.
The fix looks clean, even though it might look like there are some "object-oriented shenanigans" happening here with making _LazyKernel class a descriptor (i.e., it has a __get__ method), it makes sense to me and the comments communicate the intent clearly, so it shouldn't be too difficult to maintain in the future if needed. I think this is a good way to proceed.
Noticed the small slowdowns (~20%) in some neighborhood benchmarks, but I think that is acceptable if this PR is solving the race condition and/or lack of caching (which slows down real workflows the first time it hits but isn't necessarily appearing in the benchmark suite).
The only issue is, I still have the crash from the race condition (#1757)! Copy-pasting those lines into terminal while on this PR's branch, it still crashes locally for me within a few runs.
If you want to merge this as closing #1769 and related to 1757 but not closing 1757, I would be okay to approve. But for closing 1757 something else needs to be fixed. I'm not sure what is missing (race conditions are tricky to debug…) but I wonder if it might be as simple as moving the threading.Lock() mechanism (not sure where it would need to move to, though…)?
Collaborator
Author
|
pre-commit.ci autofix |
for more information, see https://pre-commit.ci
Sevans711
reviewed
Sep 25, 2026
Sevans711
left a comment
Sevans711
left a comment
Collaborator
There was a problem hiding this comment.
Looks like the race condition is now fixed as well! Running the code from #1757, I now see no crashes after 20 runs, whereas before it was usually crashing within a few runs at most, and never took more than 10 runs before the race condition caused a crash.
Skimmed through the code changes since last time I viewed and they look reasonable.
My only concern now is that some of the 480km Neighborhood benchmarks show slowdowns. Though, the 120km ones show speedups…? This makes me wonder if maybe the ways this PR affects when numba compilation happens here is causing the ASV benchmarking suite to not warm up properly in this case. Are these slowdowns real, or just because they are including the numba compilation times? Aside from that, I would be happy to approve!
Collaborator
Author
|
@Sevans711 Good catch, there was some numba compilation getting into the timings on smaller grids. Otherwise, there is a small overall speedup on the neighborhood reductions. Some of the other benchmark differences are probably noise (like the |
Sevans711
approved these changes
Sep 25, 2026
Sevans711
left a comment
Sevans711
left a comment
Collaborator
There was a problem hiding this comment.
Approving, now that the benchmarks look better!
Member
|
Everything seems to be working as intended and not seeing any concerns. The only thing to mention here is that writing |
dylannelson
approved these changes
Sep 29, 2026
cmdupuis3
added a commit
to cmdupuis3/uxarray
that referenced
this pull request
Sep 29, 2026
Conflict in uxarray/grid/neighbors.py: UXARRAY#1768 replaced the target="parallel" gufunc with njit kernels, running the serial one on dask blocks and _reduce_rows_parallel only on in-memory arrays. Took that version, and put _reduce_rows_parallel under @parallel_njit like the other parallel kernels: an in-memory reduction called from a user's own thread pool still aborted under workqueue. test_parallel drops the neighborhood case, whose kernel no longer exists; that path now goes through @parallel_njit like the centroid kernel it keeps. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #1757, #1769
Overview
The PR solves two issues with the current version of the neighborhood filters kernels:
njit(cache=True). This means that in the current implementation of kernel caching, every call to a kernel gets a new UUID, so any attempt to cache this way will basically be a no-op.In a fun twist, both are known unresolved issues in numba:
numba/numba#6264
numba/numba#9288
PR Checklist
General
Testing & Benchmarking
Documentation and Examples
docs/api.rst; internal (private) function names start with an underscore (_)AI Disclosure
AI Usage: Claude Opus 5