Closed shoyer closed 1 year ago
Merging #89 (847766f) into master (9d12932) will increase coverage by
0.32%
. The diff coverage is97.01%
.
@@ Coverage Diff @@
## master #89 +/- ##
==========================================
+ Coverage 96.02% 96.35% +0.32%
==========================================
Files 11 11
Lines 503 548 +45
Branches 112 105 -7
==========================================
+ Hits 483 528 +45
Misses 13 13
Partials 7 7
Impacted Files | Coverage Δ | |
---|---|---|
rechunker/api.py | 97.12% <75.00%> (-0.92%) |
:arrow_down: |
rechunker/algorithm.py | 91.22% <100.00%> (+5.51%) |
:arrow_up: |
rechunker/compat.py | 100.00% <100.00%> (+22.22%) |
:arrow_up: |
Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.
Thinking about this a little bit more, this may be a harder problem for the rechunker data model than I thought. The problem is that my algorithm does not guarantee that intermediate chunks can be written to Zarr.
For example, consider stage 1 in my proposed rechunking plan above: (93, 721, 1440) -> (93, 173, 396) -> (1447, 173, 396)
. 1447 is not a multiple of 93 -- they have a greatest common divisor of 1 -- so even though the typical chunk along the first axes in the intermediates would be 93, in the worst case we'd have to save a single element along that axis.
In Xarray-Beam and Dask, this is handled by producing temporary irregular chunks from the "split" step. This is fine in principle, but won't work for the current version of Rechunker, which stores intermediates in Zarr.
I see a few possible ways to resolve this:
Thanks a lot for this Stephan! It will take me a couple of days to digest this.
I tested this in Xarray-Beam today, rechunking a single 1.3 TB float32 array from (31, 721, 1440)
-> (350640, 10, 10)
with max_mem=256e6
. This is large enough that Rechunker doesn't do any additional consolidation of reads/writes:
Edit: you can find the example code for this over in https://github.com/google/xarray-beam/pull/14
I am happy to rebase, but would anyone be up for reviewing this PR?
The alternative would be to fork this code into Xarray-Beam.
@rabernat any thoughts here?
HI Stephan! Sorry for letting this lie dormant for so long. Will definitely take a look this week. Thanks for your patience.
HI Stephan! Sorry for letting this lie dormant for so long. Will definitely take a look this week. Thanks for your patience.
Thanks!
I just merged in "master" so this should be ready for review.
I In particular, I am not quite comprehending the what gives rise to the issue of irregular intermediate chunks, which I believe is linked to the least-common multiple calculation.
This is a good point! The code as I wrote it was rather misleading -- the LCM logic is indeed just about counting IO ops.
Irregular chunks only become necessary because of the very simple algorithm (geometric spacing) I use for picking intermediate chunk sizes in multi-stage plans. With some cleverness, this could likely be avoided in many cases.
The current rechunker algorithm does actually support "irregular chunks" in the form of overlapping reads of source arrays. So it's certainly possible to use this limited form of irregular chunking to support arbitrary rechunks, and I suspect with effort we could make it work for multi-stage rechunks, too. The most obvious way to do this would be to simply pick some intermediate chunk size, and then run the existing rechunker algorithm recursively twice, for three total temporary Zarr arrays (or 7 or 15). This would be quite a bit less flexible than what I implemented here, though.
Fantastic. These comments are a big help. Thanks for this contribution.
Background
Rechunker's current "Push / Pull Consolidated" algorithm can be thought of as a mix of "split" and "combine" steps:
source_chunks
->read_chunks
(to fill upmax_mem
per chunk)read_chunks != write_chunks
:read_chunks
->int_chunks
int_chunks
->write_chunks
write_chunks
->target_chunks
This is pretty clever, but can run into scalability issues. In particular, sometimes
int_chunks
must be very small, which results in a significant overhead for reading/writing small files (https://github.com/pangeo-data/rechunker/issues/80).Proposal
I think the right way to fix this is to extend Rechunker's algorithm to allow for multiple split/combine stages -- as many as necessary, to avoid creating tiny intermediate chunk sizes. This PR implements the math for such as algorithm, in a fully backwards compatible fashion. Users can control the number of stages via the new
min_mem
parameter, which specifies a minimum chunk size in bytes.Multi-stage rechunking is not yet hooked up to any of Rechunker's executors. I'm proposing adding it to rechunking because it reuses/modifies the existing code, and I need the math for rechunking inside Xarray-Beam (this was an easier way to explore writing the Beam executor part). Unfortunately it isn't easy to add into Rechunker's existing model, which stores intermediate results in Zarr, because multi-stage rechunking (at least this version) requires irregular chunks.
Dask also does multi-stage rechunking, which brought significant efficiency gains (https://github.com/dask/dask/issues/417). I considered copying Dask's rechunk planning algorithm here, but it involves a lot of complex logic so I decided to try replacing it with a simple heuristic instead. See
algorithm.py
for details.Example
My specific motivation for this PR is experimenting with rechunking on Pangeo's ERA5 single-level dataset, which contains a number of 1.5 TB variables (at least once decoded into float32). Rechunking these arrays with
shape=(350640, 721, 1440)
from "whole image" chunks(31, 721, 1440)
to "whole time-series" chunks(350640, 10, 10)
with Rechunker's current algorithm produces a very large number of small chunks. It works, but seems much slower than it should be.I wrote a little Python script to compare Rechunker's current method (
min_mem=0
), with my proposed multi-stage method (min_mem=10MB
) and Dask's rechunking method:Comparing my new multi-stage algorithm (
max_mem=10MB
) to Rechunker's existing algorithm (max_mem=0
), the multi-stage pipeline does two extra dataset copies, but reduces the number of IO operations by ~50x.Comparing my new algorithm to Dask's algorithm, the plans actually look remarkably similar. My estimates suggest that my algorithm should involve about half the number of IO operations, but Dask's plan uses slightly "nicer" chunk sizes. I have no idea which is better is practice, and note that I'm using Dask's algorithm without adjusting any of the control knobs.
I have not yet benchmarked any of these algorithms on real rechunking tasks.See below for benchmarking results from Xarray-Beam, for which it significantly improves real-world performance.