Zostoga still causing memory spikes

[mo-gill]

I've tested zostoga with one month of data after this fix and it seems we're still getting OOM kill events (despite onm mem limit currently being set to 100GB) so this can't be produced for the time being.

I added a bit of profiling to see where the memory is spiking:

Loading data for "CMIP7_ocean.json: zostoga_tavg-u-hm-sea"
[memory profile] calc_zostoga: start: RSS=36737.9 MB
[memory profile] calc_zostoga: before processing time step 0: RSS=63702.0 MB
[memory profile] calc_rho_mean: start: RSS=63702.0 MB
[memory profile] eos_insitu: start: RSS=63702.0 MB
[memory profile] eos_insitu: after cast to float64: RSS=93769.3 MB

I did do a quick and dirty investigation using copilot to see if there were any quick wins. Unfortunately there doesn't seem to be. It has suggested breaking things up a bit by looping over depth levels in calc_rho_mean (and some other bits), which should in theory limit the peak memory usage to ~1 levels worth of data on each loop. This solution looks like it could be quite a complex code change though so i'm not going any further. If anyone is curious as to what that suggested change might look like ( i've left the memory profiling statements in there) , you can have a look here.

Edit: I've tried running with that fix and it does produce with a maximum memory usage of ~68GB. However, i'm lacking the expertise to know if there may be a simpler solution.

Putting this on hold for the time being. Will try to fix in CDDS v4.0.1

Daley Calvert (@dcalve) any thoughts on potential solutions would be appreciated!

[dcalve]

I suspect my test case may not have replicated the CDDS workflow closely enough- it may be that cube.core_data() already contains 'real' data on entry to calc_zostoga (likely, given 36Gb of memory is in use at this point in your profiling). If this is the case, then my fixes would have no impact.

It would be useful to add a pdb.set_trace() at the point of entry to calc_zostoga and run thetao.has_lazy_data()/thkcello.has_lazy_data() to check this.

If the data payload of the cube is already loaded into memory, then I guess this is because other processors have triggered this for thetao and thkcello? Or is CDDS perhaps doing this in preparing the data? I can imagine the masking of the northfold would do this.

The only other quick solution I can think of is to break down the eos_insitu calculation of calc_rho_mean into smaller chunks, which would reduce memory use at the cost of slower processing. This might be easier than trying to retain cube.core_data() as a Dask array in CDDS.