Xarray indices and other non-Dask variables are round-tripped to the Dask scheduler
まだ誰も着手していません。
評価
- 難易度
- 5/5
- 見積もり時間
- 1週間以上
- 初心者へのやさしさ
- 35/100
- issue の種類
- バグ
- 明瞭さ
- 明確に書かれている
- 活発さ
- 停滞
- 技術スタック
- python
調査の方向性
提供された xarray reproducer を実行し、次に distributed/client.py の collections_to_expr、dask/base.py の HLGExpr.from_collection、および dask/_expr.py の参照されている処理の周辺を調べます。dask_graph、dask_keys、dask_postpersist/dask_postcompute hooks を追跡します。persist または compute の間、非 Dask collection の内容が scheduler にシリアライズされたり、scheduler 上に保存されたりしなければ完了です。
索引モデルが issue の本文から書いたものです。
説明
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
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
-
use
xarray.DataArray.persistand 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.
tmp = a.drop_vars(a.coords)
tmp, = dask.persist(tmp)
b = xarray.merge([a.coords.to_dataset(), tmp])
- 主要言語
- Python
- スター
- 1.7k
- フォーク
- 780
- PR マージ指標
- 30日以内にマージされた PR はありません
環境構築
- Dockerfile・Docker Compose ファイルなし
- プルリクエストのテンプレートあり
- コントリビューションガイドを読む
はじめの一歩
- issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
- 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
- リポジトリをフォークし、ブランチを切って変更します。
- issue 番号を参照したプルリクエストを送ります。
dask/distributed のほかの issue
-
needs triage
難易度 2/5 1〜3時間 初心者へのやさしさ 72/100
dask/distributed#9366 ·
-
needs triage
難易度 2/5 1〜3時間 初心者へのやさしさ 84/100
dask/distributed#9353 ·
-
documentation
難易度 1/5 1〜3時間 初心者へのやさしさ 82/100
dask/distributed#8304 ·
-
難易度 2/5 1〜3時間 初心者へのやさしさ 74/100
dask/distributed#4816 · コメント 2 件 ·
-
documentation good first issue
難易度 2/5 1〜3時間 初心者へのやさしさ 74/100
dask/distributed#2378 · コメント 2 件 ·
dask/distributed の issue をすべて見る
似ている issue
-
bug server
難易度 2/5 1〜3時間 初心者へのやさしさ 78/100
メンテナーはふだん 1 日以内に返信
-
難易度 2/5 1〜3時間 初心者へのやさしさ 78/100
sportsdataverse/sportsdataverse-py#641 ·
メンテナーはふだん 1 日以内に返信
-
難易度 2/5 1〜3時間 初心者へのやさしさ 82/100
googleapis/google-cloud-python#18532 ·
メンテナーはふだん 1 日以内に返信
-
難易度 2/5 1〜3時間 初心者へのやさしさ 72/100
メンテナーはふだん 1 日以内に返信
-
難易度 2/5 1〜3時間 初心者へのやさしさ 72/100
メンテナーはふだん 1 日以内に返信