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

Dask Array.compute() peak memory in Jupyterlab

我在分布式集群上使用 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

当你这样做时会发生什么 compute():

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

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