使用 Zarr 存储 Dask 数组消耗太多内存

Question

我有一长串 .zarr 数组，我想将它们合并为一个数组并写入磁盘。

我的代码大致如下：

import dask.array
import zarr
import os

local_paths = ['parts/X_00000000.zarr',
 'parts/X_00000001.zarr',
 'parts/X_00000002.zarr',
 'parts/X_00000003.zarr',
 'parts/X_00000004.zarr',
 'parts/X_00000005.zarr',
 ...]

result_path = "testtest"
os.makedirs(result_path)

Xs = [dask.array.from_zarr(zarr.DirectoryStore(p)) for p in local_paths]
X = dask.array.concatenate(Xs, axis=0)
X = X.rechunk({0: 10000, 1:-1, 2:-1, 3:-1})
dask.array.to_zarr(X, zarr.DirectoryStore(result_path))

local_paths 中的每个数组都包含一个 64x64 图像列表。这些列表的长度各不相同。所以第一个的形状可能是 (100, 64, 64, 3)，第二个的形状可能是 (200, 64, 64, 3)。

执行这段代码的最后一行，导致我的内存完全耗尽，然后 Jupyter notebook 完全崩溃（没有给我错误消息或异常）。

为了调查问题，我打印了任务图，因此用以下两行替换了最后一行：

k = dask.array.to_zarr(X, zarr.DirectoryStore(result_path), compute=False)
k.visualize()

它非常大（link），所以我只截取了其中两个有趣的部分：

这个结构一直重复。 Dask 获取连接的输出，重新分配数据，然后尝试存储它。请注意由重叠过渡产生的粗黑条。

现在看看这些转换是从哪里来的：

看中间的 create 节点。我假设这是图中创建 zarr DirectoryStore 的部分。 create 节点的前驱节点已过渡到所有存储节点！

这是我猜测为什么 dask 内存不足。它首先尝试解决所有重新分块合并问题，到它应该创建 DirectoryStore 时，已经没有内存了。并且所有 store 节点都无法执行，因为 create 节点是每个节点的前提条件。

我的假设是真的吗？如果是，我该怎么做才能强制 dask 首先创建 DirectoryStore？如果没有，还有什么可能是我内存不足的问题？

更新 当我使用 dask.config.set(scheduler='single-threaded') 时，DirectoryStore（创建节点）的创建不是问题。我刚刚查看了输出目录，已经写入了一些文件。所以一定是任务图本身太大了。