问同时将xarray数据集写入zarr --如何有效地使用dask分布式扩展

EN

TLDR:

如何有效地使用dask-distributed将许多dask-backed xarray数据集写入AWS S3上的zarr存储区？

详细信息

我有一个工作流程，它在S3上获取栅格数据集的列表，并生成一个dask数组支持的xarray数据集。

我需要迭代多个组，其中每个组的工作流获取属于该组的栅格数据集，并生成相应的xarray数据集。

现在，我希望将数据集中的数据写入S3上的扎尔存储(同一存储区，仅使用group参数)。

顺序处理的伪代码是这样的：

client = Client(...) # using a distributed cluster

zarr_store = fsspec.get_mapper("s3://bucket/key.zarr")

for group_select in groups:
    
    xr_dataset = get_dataset_for_group(group_select)
    
    # totally unnecessary, just to illustrate that this is a lazy dataset, nothing has been loaded yet
    assert dask.is_dask_collection(xr_dataset)
    
    xr_dataset.to_zarr(zarr_store, group=group_select)

这很好，一旦执行to_zarr，数据就被加载并存储在S3上，任务并行运行。

现在，我想使用dask.distribuited并行运行它。这是我尝试过的，也是我遇到的问题：

1.使用.to_zarr(..., compute=False)收集延迟任务列表的

这在原则上是可行的，但相当缓慢。创建一个任务需要3到4秒左右，我需要运行这个100+时间，在任何计算真正开始之前需要4到5分钟。

2.将其封装到dask.delayed中

这极大地加快了任务的创建速度，但是给zarr存储的写入并不是在工人之间分开，而是处理任务的工作人员在完成加载任务并将其写入zarr之后收集所有数据。

3.用自定义函数包装to_zarr并将其传递给client.submit

这似乎是最有希望的选择。我刚刚将to_zarr调用包装在一个自定义函数中，该函数可以从一个工作人员调用：

def dump(ds, target, group=None):
    with worker_client() as client:
        ds.to_zarr(store=target, group=group)  
    return True

使用worker_client执行此操作将编写任务返回到调度程序，并解决了我在dask.delayed中遇到的问题。

但是，当我重复提交此函数时(我需要按以下方式执行此100+时间)

futures = [client.submit(dump, x, target, g) for x,g in zip(datasets, groups)]

我很快就用要处理的任务压倒了调度程序。

对于这一点，我能想到的唯一明显的解决方案是将数据集分成几个批次，并且只有在前一个数据集完成后才启动一个新的数据集。但难道没有更优雅的解决方案吗？或者在dask (分布式)中有内置功能吗？