Question

我有一个包含4个工作程序的Dask-MPI集群，一个3D网格数据集已加载到Dask数组中，并分块为4个块。我的应用程序要求我为每个工作人员准确地运行一个任务，最好每个任务运行一个块。我遇到的麻烦是以可靠，可复制的方式使块分布在整个群集中。具体来说，如果我运行array.map_blocks（foo），则foo在每个块的同一工作线程上运行。

Client.rebalance（）似乎应该执行我想要的操作，但是它仍然将所有或大多数块留在同一工作线程上。作为测试，我尝试将数据重新分块为128个块并再次运行，这导致7或8个块移至其他数据集。这暗示Dask正在使用一种启发式方法来决定何时自动移动块，但是并没有给我一种强制均匀分布块的方法。

这是我一直在尝试的简单测试脚本（连接到具有4个工作人员/等级的集群）。

#connect to the Dask scheduler
from dask.distributed import Client, Sub, Pub, fire_and_forget
client = Client(scheduler_file='../scheduler.json', set_as_default=True)


#load data into a numpy array
import numpy as np
npvol = np.array(np.fromfile('/home/nleaf/data/RegGrid/Vorts_t50_128x128x128_f32.raw', dtype=np.float32))
npvol = npvol.reshape([128,128,128])

#convert numpy array to a dask array
import dask.array as da
ar = da.from_array(npvol).rechunk([npvol.shape[0], npvol.shape[1], npvol.shape[2]/N])


def test(ar):
    from mpi4py import MPI
    rank = MPI.COMM_WORLD.Get_rank()
    return np.array([rank], ndmin=3, dtype=np.int)

client.rebalance()
print(client.persist(ar.map_blocks(test, chunks=(1,1,1))).compute())

在几十次测试运行中，这段代码一次返回了3级的一个块，否则将所有的块都置于0级。

Answer 1

由于总数据集并不大，因此对from_array的初始调用仅创建了一个块，因此将其分配给一个工作程序（您可能用chunks=指定了）。以下rechunk尽可能不移动数据。

假设每个工作人员都可以访问您的文件，那么最好首先将工作块中的数据块卸载。

您将需要类似的功能

def get_chunk(fn, offset, count, shape, dtype):
    with open(fn, 'rb') as f:
        f.seek(offset)
        return np.fromfile(f, dtype=dtype, count=count).reshape(shape)

，并为每个块传递不同的偏移量。

parts = [da.from_delayed(dask.delayed(get_chunk)(fn, offset, count, shape, dtype), shape, dtype) for offset in [...]]
arr = da.concat(parts)

这与npy source的Intake代码https://github.com/intake/intake/blob/master/intake/source/npy.py#L11

中自动完成的操作非常相似。

使用dask.distributed强制或显式重新平衡数据

1 个答案: