Implement .to_chunk_cached_arrays() with ChunkCachedArray - #2860
Implement .to_chunk_cached_arrays() with ChunkCachedArray#2860VeckoTheGecko wants to merge 13 commits into
.to_chunk_cached_arrays() with ChunkCachedArray#2860Conversation
Provides an opt-in optimization that wraps dask-backed field data in chunk-level LRU caches via the chunk_cached_array package. Repeated vectorized .isel() calls hit an in-memory cache instead of recomputing dask task graphs, giving large speedups for particle simulations.
Results on ds_2d_left_agrid.zarr with 10k particles: - Plain dask: 85.0s (1x) - Windowed arrays: 16.8s (5.1x) - Cached chunk arrays: 4.9s (17.3x)
Moves the code to a Parcels subpackage (rather than a separate package).
This was a testing-only convenience method, now handled internally by Xarray
instead of `cached_chunk`
Thanks for this exciting PR, @VeckoTheGecko! I plan to do some real-world performance testing (also including/comparing #2846 and v3) later this week. Will report back here when I know more! |
|
This is really cool! The potential it has to lower peak memory usage for comparable speed to the current windowed array implementation is very exciting. I left a couple of comments and will also plan on doing some performance testing through the week! |
I can't see them on my end, maybe you still need to click 'submit' on the review? |
|
|
||
| def put(self, key: Hashable, value: np.ndarray) -> None: | ||
| nbytes = value.nbytes | ||
| if nbytes > self._max_bytes: |
There was a problem hiding this comment.
I wonder if this should be accompanied by a warning. My understanding is the chunk size is set by dask/xarray and if it is greater than the cache size then this would silently skip use of the cache altogether.
There was a problem hiding this comment.
I think that's reasonable to raise a warning. Any thoughts about how we can only raise it once? (since raising on every put would flood the output)
| flat_keys = np.ravel_multi_index(chunk_ids, numblocks) | ||
|
|
||
| # Sort points by flat chunk key to group them. | ||
| sort_order = np.argsort(flat_keys, kind="mergesort") |
There was a problem hiding this comment.
From the testing I've done it looks like for randomly distributed particles quicksort may be faster here?
There was a problem hiding this comment.
Awesome! I'll update
3 participants
Description
Provides an opt-in optimization that wraps dask-backed field data in
chunk-level LRU caches.
vectorized .isel() calls hit an in-memory cache instead of recomputing
dask task graphs, giving large speedups for particle simulations.
This PR adds
ChunkCachedArrayaccording to the approach detailed in #2854 , providing an opt-in optimization wrapping dash-backed field data with (dask)chunk-level LRU caches on the array level. This PR:parcels._chunk_cached_arraywith the implementation of the Array and the needed LRU cache. Note the cache is on the array level with configurable memory limits per array.parcels._chunk_cached_array.wrap_datasetproviding the main entrypoint for Parcels for wrapping Xarray dataset objects.to_chunk_cached_arraystoFieldSetandModelDataclasses.test_backends.py) on a NEMO dataset, asserting that the same results are received regardless of backend.Still needed:
data-generation.py,benchmark_chunk_cache.py,benchmark_chunk_cache.png)Future work:
max_cache_bytesonFieldSet.to_chunk_cached_arraysoperates per variable (i.e., perField). This requires the user to think about how many Fields are in their simulation in order to tune their performance. Down the line, we could tune this to allocate a total cache of (e.g.,) 50% of the system memory (which would allow for other memory overhead of coordinates (not stored in the cache), the Python process, or other user processes). This would require significant refactoring not only on how the chunk cache is handled, but also backends in general.(cc @erikvansebille, @wyatt-fluidnumerics )
Performance
I have some some preliminary profiling with a 20Gb idealised dataset, and found the following performance profile which looks like a promising first step.
I've included
data-generation.pyandbenchmark_chunk_cache.pyhere for testing in case it helps.Its not clear to me the performance with real-world simulations and how that scales (I only have access to my laptop, and found working/debugging on Lorenz quite frustrating with my high latency from travels). Let me know if there's anything I can do to help here @erikvansebille .
Checklist
ChunkCachedArray#2854mainfor normal development,v3-supportfor v3 support)AI Disclosure