dask / dask/distributed

Xarray indices and other non-Dask variables are round-tripped to the Dask scheduler

Open
#9,200 6 comments 0 reactions 0 assignees View on GitHub
memory p2 performance
Dominant language
Python
Stars
1.7k
Forks
778
Avg merge
2h 50m
Merged PRs (30d)
3

Description

Upon calling one of
- `dask.persist()`
- `dask.compute()`
- `distributed.Client.persist()`
- `distributed.Client.compute()`

on a `xarray.DataArray` or `xarray.Dataset`, the indices and all other non-Dask variables are serialized and sent to the Dask scheduler and then back to the client when the computation is done. This is unnecessary and can be extremely costly, as the whole data is serialized, sent over the network, and stored on the scheduler for the whole duration of the computation. In case of object string variables, there is also a substantial impact on the GIL.

This issue also impacts any other third-party library that defines Dask collections with non-trivial contents outside of their Dask graph.

## Reproducer

```python
import dask
import dask.array as da
import numpy as np
import distributed
import xarray
import pickle
from dask.utils import format_bytes

if __name__ == "__main__":
a = xarray.DataArray(
da.random.random(10_000_000),
dims=["x"],
coords={"x": np.arange(0, 100_000_000, 10)},
)

print("__dask_graph__()", format_bytes(len(pickle.dumps(a.__dask_graph__()))))
print("__dask_keys__()", format_bytes(len(pickle.dumps(a.__dask_keys__()))))
print("__dask_postpersist__()", format_bytes(len(pickle.dumps(a.__dask_postpersist__()))))
print("__dask_postcompute__()", format_bytes(len(pickle.dumps(a.__dask_postcompute__()))))

with distributed.Client() as client:
b = client.persist(a) # sends 152 MiB
# b, = dask.persist(a) # sends 152 MiB
# b = a.persist() # sends 3 kiB; see below
c = b.sum()
del b
c.compute()
```
Output:
```
__dask_graph__() 3.10 kiB
__dask_keys__() 71 B
__dask_postpersist__() 152.59 MiB
__dask_postcompute__() 152.59 MiB
/home/crusaderky/github/enkiext/.pixi/envs/default/lib/python3.14/site-packages/distributed/client.py:3387: UserWarning: Sending large graph of size 152.59 MiB.
This may cause some slowdown.
Consider loading the data with Dask directly
or using futures or delayed objects to embed the data into the graph without repetition.
See also https://docs.dask.org/en/stable/best-practices.html#load-data-with-dask for more information.
```

## Analysis
Upon calling `persist()`, the Dask Client should dismember the xarray collection by calling its hooks `__dask_graph__()`, `__dask_keys__()`, and `__dask_postpersist__()`. Upon calling `compute()`, Dask should instead use `__dask_graph__()`, `__dask_keys__()`, and `__dask_postcompute__()`.

Crucially, in xarray the return value of `__dask_postpersist__()` and `__dask_postcompute__()` embeds all non-Dask variables. The output of these functions is supposed to be stored on the Client and never serialized.

Instead, `distributed.Client.persist` calls `collections_to_expr`:
https://github.com/dask/distributed/blob/991fda7cf985289631a2590aac422ffb8c0d4adf/distributed/client.py#L3811

which in turn calls `HLGExpr.from_collection`:
https://github.com/dask/dask/blob/714eaf4e3515ec5c65cbc61af22a932c84d156c6/dask/base.py#L445

which in turn embeds the whole contents of `__dask_postpersist__` in its return value, which is then sent to the scheduler:
https://github.com/dask/dask/blob/714eaf4e3515ec5c65cbc61af22a932c84d156c6/dask/_expr.py#L1020

## Workarounds
1. use `xarray.DataArray.persist` and other xarray methods, which are unaffected due to the apparently convoluted way they are implemented: https://github.com/pydata/xarray/blob/37f2d49b5cfbf5ca7e24dbad6347f99f5a24a368/xarray/core/dataset.py#L800-L815
This however can result in inefficiency when persisting, and even more so computing, multiple collections at once.

2.
```python
tmp = a.drop_vars(a.coords)
tmp, = dask.persist(tmp)
b = xarray.merge([a.coords.to_dataset(), tmp])
```

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.