A lot of computational memory and time is being devoted to reshaping and rechunking

[dhjx1996]

Describe the bug
A lot of computational memory and time seems to be devoted to reshaping and rechunking in compute_cloud_optics (and possibly in compute_gas_optics -- untested). This is likely related to the inconsistency between the chunking of the input and output (#191).

To Reproduce
Steps to reproduce the behavior

1. Load atmosphere of type xarray.Dataset:

1. cloud_optics = load_cloud_optics(cloud_optics_file=CloudOpticsFiles.SW_BND)
2. clouds_optical_props = cloud_optics.compute_cloud_optics(atmosphere, problem_type="two-stream")
3. I ran the following code to monitor progress, but the essential line is clouds_optical_props = clouds_optical_props.persist()

# Progress bar with the distributed scheduler
from dask.distributed import Client, progress

client = Client()  # use dask.distributed by default

clouds_optical_props = clouds_optical_props.persist()  # start computation in the background
progress(clouds_optical_props )      # watch progress

Here is the progress report of the above block of code:

I observed server memory spiking to around 80gb during the computations and my server eventually crashes (the server memory limit is 128gb so maybe a memory overflow caused the crash but I am not fully certain as to what causes the crash). As an alternative to step 3, I have tried tau_arr_swbndsum = clouds_optical_props.tau.sum("bnd", skipna=True).persist() (and analogously for g and ssa) and that code works but I still observe memory spikes to around 50gb. The progress report:

Expected behavior
Is all that rechunking and reshaping necessary?

Software Versions:

• OS: Ubuntu 22.04.5 LTS
• pyrte_rrtmgp, rte_rrtmgp, dask, xarray versions: 0.0.9, 1.9.1, 2025.3.1, 2025.3.0

[brendancol]

@dhjx1996 I'm confident that changes are needed to pyRTE to get the dask.distributed case "purrrring", so thank you for your patience in testing as I'm sure you've invested some time here.

Expected behavior
Is all that rechunking and reshaping necessary?

The number of rechunking and reshaping / shuffling calls is proportional to the size of the task graph.

It looks like the chunksize is a bit small at 12MB for atmosphere in the screenshot above. As a rule of thumb, try partitions between 512MB-2GB (.rechunk()) depending on your exact worker profile...something around there which has a clean shape.

This will make the task graph smaller, thus fewer operations. You could then also remove the persist() and end with a to_zarr() call to write result chunks to disk...or reduce / resample results into an image that fits on the client....anything doesn't bring the full result back to a single node or have too much cross-worker serialization.

Good next steps here for pyRTE would be:

• Run memory_profiler to find leaks and unnecessary allocations
• Run line_profiler to understand line-by-line relative performance

[dhjx1996]

@brendancol Thanks for the reply! My understanding is that if persist() does not work compute() will not either. I have tried cloud_optical_props.to_zarr(compute=True) (the default) but that gives memory problems. I also tried .to_zarr(compute=False) and that works but when I load the data and call .max(skipna=True).values on any of the variables it returns 0 which is clearly incorrect.

I tried a much larger chunk size and here are the results.

Dataset

Progress

My chunk size is 814.75mb (is that per array though?) and the number of computations went down a lot! However, my workers are now running out of memory. I understand that Dask generally has a trade-off between number of computations and memory usage (chunk size), in that case I am slightly skeptical that increasing chunk size will fix the problem if the initial problem, despite the tiny chunks, was memory. I will continue experimenting with other chunk sizes and report back!