Task graphs on .map_blocks with many chunks can be huge

What happened?

I'm getting task graphs > 1GB, I think possibly because the full indexes are being included in every task?

What did you expect to happen?

Only the relevant sections of the index would be included

Minimal Complete Verifiable Example

da = xr.tutorial.load_dataset('air_temperature')

# Dropping the index doesn't generally matter that much...

len(cloudpickle.dumps(da.chunk(lat=1, lon=1)))
# 15569320

len(cloudpickle.dumps(da.chunk().drop_vars(da.indexes)))
# 15477313

# But with `.map_blocks`, it really matters — it's really big with the indexes, and the same size without:


len(cloudpickle.dumps(da.chunk(lat=1, lon=1).map_blocks(lambda x: x)))
# 79307120

len(cloudpickle.dumps(da.chunk(lat=1, lon=1).drop_vars(da.indexes).map_blocks(lambda x: x)))
# 16016173

MVCE confirmation

Minimal example — the example is as focused as reasonably possible to demonstrate the underlying issue in xarray.
Complete example — the example is self-contained, including all data and the text of any traceback.
Verifiable example — the example copy & pastes into an IPython prompt or Binder notebook, returning the result.
New issue — a search of GitHub Issues suggests this is not a duplicate.
Recent environment — the issue occurs with the latest version of xarray and its dependencies.

Relevant log output

No response

Anything else we need to know?

No response

Environment

Details

INSTALLED VERSIONS

commit: None
python: 3.9.18 (main, Aug 24 2023, 21:19:58)
[Clang 14.0.3 (clang-1403.0.22.14.1)]
python-bits: 64
OS: Darwin
OS-release: 22.6.0
machine: arm64
processor: arm
byteorder: little
LC_ALL: en_US.UTF-8
LANG: None
LOCALE: ('en_US', 'UTF-8')
libhdf5: 1.12.2
libnetcdf: None

xarray: 2023.10.1
pandas: 2.1.1
numpy: 1.26.1
scipy: 1.11.1
netCDF4: None
pydap: None
h5netcdf: 1.1.0
h5py: 3.8.0
Nio: None
zarr: 2.16.0
cftime: 1.6.2
nc_time_axis: None
PseudoNetCDF: None
iris: None
bottleneck: 1.3.7
dask: 2023.5.0
distributed: 2023.5.0
matplotlib: 3.6.0
cartopy: None
seaborn: 0.12.2
numbagg: 0.6.0
fsspec: 2022.8.2
cupy: None
pint: 0.22
sparse: 0.14.0
flox: 0.7.2
numpy_groupies: 0.9.22
setuptools: 68.1.2
pip: 23.2.1
conda: None
pytest: 7.4.0
mypy: 1.6.1
IPython: 8.14.0
sphinx: 5.2.1

They should be getting subset here:

xarray/xarray/core/parallel.py

Line 421 in 83fbcf0

def subset_dataset_to_block(

It's possible that index vars are getting duplicated a bunch, since they're "broadcast" out to every chunk. We might be able to be clever here.

Yes that's exactly what's happening 🤦🏾‍♂️

da = xr.tutorial.load_dataset("air_temperature")
chunked = da.chunk(lat=10, lon=10)
chunked.chunksizes
# Frozen({'time': (2920,), 'lat': (10, 10, 5), 'lon': (10, 10, 10, 10, 10, 3)})

mapped = chunked.map_blocks(lambda x: x)

chunked.air.data.numblocks 
# (1,3,6) = 18 blocks in total

len([key for key in dict(mapped.__dask_graph__()) if "time" in key[0]]) 
# 18!

The time vector is included in the graph 18 times, even though there is only a single chunk along the time dimensions

Minimize duplication in `map_blocks` task graph

The other piece of duplication is here:

xarray/xarray/core/parallel.py

Lines 497 to 500 in 83fbcf0

    
           expected["indexes"] = { 
        
               dim: indexes[dim][_get_chunk_slicer(dim, chunk_index, output_chunk_bounds)] 
        
               for dim in indexes 
        
           }

If I comment that out, I get close to the expected size with no indexes

Quick question re a workaround, if anyone knows off-hand — would creating the job from a delayed task get around this?

My (low-confidence) understanding is that there are two downsides of big task graphs:

A task graph that's huge in number of tasks, even though each task is small. This slows the scheduler down because it has to process the graph.
A task graph that contains big objects, even though there aren't many tasks. This requires sending big objects to the scheduler. But possibly if the job is created from a worker, then the objects are already tracked by the scheduler, and it doesn't need to transfer them?

I've generally had a lot of trouble with the first, so have been trying to increase chunk size and break jobs into multiple sections (and using rechunker to avoid rechunking within a job, thanks rechunker!). But is this a case of the 2nd?

A bit late to the conversation here but ... I've had some success in this arena by telling dask to use the threaded or synchronous scheduler within the map-blocks function. It ends up looking like this:

def my_func(da: xr.DataArray) -> xr.DataArray:
    with dask.config.set(scheduler='single-threaded'):
        da2 = da - da.mean(dim='time')
    return da2

da2 = da.map_blocks(my_func)

I'm seem to remember folks telling me this was a bad idea but it made sense for my applications where I wanted to think about my arrays inside my function as numpy-arrays, not dask-arrays.

#8412 is the right solution here IMO.

Minimize duplication in `map_blocks` task graph

Minimize duplication in `map_blocks` task graph (#8412)

Share map_blocks coordinate tasks across calls

Navigation Menu

Repository navigation

Task graphs on `.map_blocks` with many chunks can be huge #8409

What happened?

What did you expect to happen?

Minimal Complete Verifiable Example

MVCE confirmation

Relevant log output

Anything else we need to know?

Environment

INSTALLED VERSIONS

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	expected["indexes"] = {
	dim: indexes[dim][_get_chunk_slicer(dim, chunk_index, output_chunk_bounds)]
	for dim in indexes
	}

Task graphs on .map_blocks with many chunks can be huge #8409

Description

What happened?

What did you expect to happen?

Minimal Complete Verifiable Example

MVCE confirmation

Relevant log output

Anything else we need to know?

Environment

INSTALLED VERSIONS

Activity

dcherian commented on Nov 3, 2023

dcherian commented on Nov 3, 2023

dcherian commented on Nov 3, 2023

max-sixty commented on Dec 8, 2023

jhamman commented on Dec 8, 2023

dcherian commented on Dec 15, 2023

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

Task graphs on `.map_blocks` with many chunks can be huge #8409