用于简单计算的大量 dask 任务

Extremely large number of dask tasks for simple computation

import numpy as np
import dask.array as dda

def index_nearest(to_find, coord_grid):
    """Return indexes of elements in coord_grid that most 
    closely match the ones in to_find.
    """
    to_find = to_find.reshape(-1,1)
    return np.argmin(np.abs(to_find - coord_grid), axis=1)

n_points = 10000
n_grid = 800

x = np.random.uniform(0, n_grid, [n_points,])
x_grid = np.arange(n_grid)

x_da = dda.from_array(x, chunks=10)
x_grid_da = dda.from_array(x_grid, chunks=10)

index_nearest(x_da, x_grid_da)

当 运行 最后一行时,我收到一条警告,提示块数已增加 40 倍,并且 repr 如下所示:

270,000 个任务对于如此简单的计算来说是不是有点太多了?

任务的数量取决于块的数量,特别是对于需要不同块的成对组合的操作。

由于具有 10 个元素的块占用 80 个字节,因此通常有足够的空间来使用更大的块,例如将块大小增加到 1_000 个元素会将单个块的内存负载增加到大约 8 KB,而只需创建 x_da 的任务数量从 1000 减少到 10。使用块大小 1_000 将完整操作的任务数减少到 61 个任务。

1_000 的块大小仍然相当小,在许多情况下,人们可能可以在 100 MB 或更大规模的块大小上逃脱(取决于硬件和执行的计算类型) .