Multi-stage rechunk

[tomwhite]

The rechunker algorithm has support for multi-stage rechunking, which reduces the number of IO operations.

It would be good to enable this in Cubed rechunk. The simplest way to do this would be to expose the min_mem parameter, but it might be interesting to see if there is a way of automatically finding the number of stages that optimizes for something we care about (e.g. min IO ops).

[tomwhite]

I rechunked one variable (1.5TB) of the ERA5 dataset in 8m 48s running on AWS Lambda using the code in cubed-dev/cubed-benchmarks#29.

The changes that were important were (see #700):

• Understanding the extra memory required by running memray on S3 Zarr code: https://github.com/tomwhite/memray-array
• Changing the rechunking algorithm to use a 7x multiplier to take account of all the copies needed
• Using a runtime image of 3.5GB (increased from 2GB) to provide enough memory (rather than using less at the expense of more rechunk stages).

[dcherian]

Amazing!

[tomwhite]

For the record, here are the raw event stats for this run:

compute-20250312T161522.549273.zip

And here are the (regenerated) timeline and memory usage (see #734) visualizations: