groupby + max() for large 2d arrays giving 'Kernel died' message

[tlsw231]

What I am trying to do: Compute groupby('time.year').max() for a 2d array with dimensions (time=237124, ncol=777602)

What is the problem: Xarray just freezes for this size.

The MVE is towards the end of this thread.

I tried using map_blocks with a function that takes a grouped array (by year in this case) as input and just returns .max('time') as the output. This works for small sizes but for large sizes again just appears to freeze [possibly, the tasks might start executing after a long wait as in some of my other map_blocks workflows].

This also brings up another issue related to the answer for a question in a previous thread where my MVE was trying to use map_blocks and a commenter pointed out that it is unnecessary for a calculation that dask should be parallelizing anyway. My understanding (could be wrong) is that there are things dask could do more efficiently but doesn't. Which is why something like map_blocks is even necessary. This always confuses me as I look at most map_blocks examples and wonder why isn't dask already figuring out all this on its own? [I realize this might be a naive question.]

MVE:

# Import various modules
import xarray as xr              
import numpy as np             
import cftime
import dask
import dask.array as da
from dask.distributed import Client
from dask_jobqueue import SLURMCluster

dask.config.set(**{'array.slicing.split_large_chunks': True})

# Request cores
ncores=20; nprocs=ncores; totalcores=ncores
cluster = SLURMCluster(cores=ncores, processes=nprocs, memory="384GB",queue="short",
                       walltime="01:00:00")
cluster.scale(totalcores)
client = Client(cluster)

# Define array
ncol=100000; # execution time < 1 min
#ncol=700000; # Kernel dies for this size which is closer to my original problem

# The duration and frequency closely match those in my original problem
start_yr = 1938; end_yr = 2100
nyears = end_yr - start_yr + 1
times=xr.cftime_range(start=format(start_yr,'04d')+'-01',
                      end=format(end_yr,'04d')+'-12-31',freq='6H',calendar='noleap')
nt=len(times)  # = 237977

data = da.random.random([nt,ncol])

dummy = xr.DataArray(data,dims=['time','ncol'],
                    coords=dict(time=(['time'],times)) ).chunk('auto')

# Compute maximum over each year for the entire 1938-2100 timeseries
dummy_annual_max = dummy.groupby('time.year').max().compute()

[tlsw231]

I know one way to solve the problem is just to split the ncol dimension into successive intervals, run several instances of the code simultaneously working on each interval and stitch the individual outputs together. Not the most elegant solution but I did this for another problem where I had to deal with a global climate dataset (i.e., large ncol). But I would like to know if there is a better alternative.

[tlsw231]

Just upgraded dask from 2021.7.0 to 2022.7.1 from and the problem completely vanished.

[dcherian]

In general you should chunk when the data is created, like with da.random.random or open_dataset.

I'm also curious if installing flox and upgrading to latest xarray makes an improvement