delayed function applied to an Array duplicates the whole dask graph on creation
- Dominant language
- Python
- Stars
- 13.9k
- Forks
- 2k
- PR merge metrics
- No merged PRs in 30d
Description
Consider this seemingly innocent snippet:
```python
a = xarray.open_zarr("foo.zarr")
b = a + 1 # Stand-in for whatever elaboration of a there may be
total = b.sum()
write_delayed = b.to_zarr("b.zarr", mode="w", compute=False)
total, _ = dask.compute(total, write_delayed)
```
This is a pretty realistic minimal use case: a user wants to write a result to zarr _and_ wants to retrieve some aggregated insight about it to the client.
Alternatively, the user may want to do store multiple outputs from the same input:
```python
a = xarray.open_zarr("foo.zarr")
b = a + 1
c = a * 2
dask.compute(
b.to_zarr("b.zarr", mode="w", compute=False),
c.to_zarr("c.zarr", mode="w", compute=False),
)
```
*The above snippets read the whole input file from disk twice.*
## What is happening
`xarray.Dataset.to_zarr(..., compute=False)` calls `dask.delayed` to collect the completion of all chunk writes:
https://github.com/pydata/xarray/blob/37f2d49b5cfbf5ca7e24dbad6347f99f5a24a368/xarray/backends/writers.py#L805-L806
This in turn triggers these lines:
https://github.com/dask/dask/blob/8a9b92f2f19d1da0172c1f6fdae0cb55dda399ae/dask/delayed.py#L179-L180
and ProhibitReuse delivers on what it is meant to do, which is duplicate the whole graph:
https://github.com/dask/dask/blob/8a9b92f2f19d1da0172c1f6fdae0cb55dda399ae/dask/_expr.py#L1353-L1357
## Original rationale
This was introduced in https://github.com/dask/dask/pull/11881. In the original comment, @fjetter mentions
> We are ensuring that whatever result comes out of this optimization is unique and the results are not being reused. This is a common problem we've already encountered in dask-expr which can cause major memory blow up otherwise.
I think that he refers to the long-standing problem where you have multiple calls to persist() that optimize differently the same segments of graph (I did not dig through the issues to find the original ticket). So he decided to duplicate the whole graph on the line before he optimized it.
The problem here is that the optimization happens halfway through the graph definition, instead of strictly on a call to `persist`, `compute`, or `optimize`.
## Affected API
- This issue impacts any delayed function that directly consumes a da.Array as an input parameter. This is the case of `xarray.Dataset.to_zarr(..., compute=False)`.
- It does NOT impact just plain `da.Array.to_zarr(..., compute=False)`, which returns a dummy da.Array instead.
- It does NOT impact a delayed function that consumes another Delayed.
- It does NOT impact `da.Array.to_delayed`
## Workaround
Notably, the last two points above allow for a workaround.
instead of calling:
```python
x: da.Array
y = dask.delayed(f)(x)
```
the issue disappears if you call
```python
x: da.Array
x_delayed = x.rechunk(-1).to_delayed().reshape(-1)[0]
y = dask.delayed(f)(x_delayed)
```
which is functionally identical, but doesn't trigger optimization.
## Alternative workaround
On the distributed scheduler, the alternative is to call `persist` before forking the graph, e.g. change
```python
a = xarray.open_zarr("foo.zarr")
total = a.sum()
write_delayed = a.to_zarr("b.zarr", mode="w", compute=False)
total, _ = dask.compute(total, write_delayed)
```
to
```python
a = xarray.open_zarr("foo.zarr")
total = a.sum()
a, total = dask.persist(a, total)
write_delayed = a.to_zarr("b.zarr", mode="w", compute=False)
del a
total, _ = dask.compute(total, write_delayed)
```
The above pattern however can be very problematic for novice users, as they risk not realising that, unless they call `del a` before blocking, they will cause the entire contents of `a` to accumulate in cluster memory, until the function calling this snippet returns.
Contributor guide
Assessment
This issue has not been assessed yet.