Jupyterlab 中的 Dask Array.compute() 峰值内存

Question

我正在分布式集群上使用 dask，当将结果返回到本地进程时，我注意到内存消耗达到峰值。

我的最小示例包括实例化集群并使用 dask.array.arange 创建一个约 1.6G 的简单数组。

我预计内存消耗约为数组大小，但我观察到内存峰值约为 3.2G。

Dask在计算过程中是否有任何副本？或者Jupyterlab需要制作一份副本吗？

import dask.array
import dask_jobqueue
import distributed

cluster_conf = {
    "cores": 1,
    "log_directory": "/work/scratch/chevrir/dask-workspace",
    "walltime": '06:00:00',
    "memory": "5GB"
}

cluster = dask_jobqueue.PBSCluster(**cluster_conf)
cluster.scale(n=1)
client = distributed.Client(cluster)
client

# 1.6 G in memory
a = dask.array.arange(2e8)

%load_ext memory_profiler
%memit a.compute()
# peak memory: 3219.02 MiB, increment: 3064.36 MiB

Answer 1

当你执行compute()时会发生什么：

• 构建计算图（这很小）并将其发送到调度程序，
• 调度程序让工作人员生成数组的各个部分，这应该是一个总数在工作人员上大约 1.6GB，
• 客户端为您要求的输出构造一个空数组，知道其类型和大小，
• 客户端通过网络或 IPC 从每个具有输出片段的工作人员接收字节串。这些复制到客户端的输出中，
• 完整的数组返回给您

您可以看到，这里的倒数第二步必然需要复制数据。原始字节缓冲区最终可能会在稍后被垃圾收集。