有没有办法限制函数调用的内存使用？

Question

我正在尝试执行一些不适合 GPU 的代码（这也发生在我的 CPU 内存中，我们的数据通常存储为 zarr 数组），但我不确定如何使用黄昏。

我发现了这个example，我正在遵循类似的策略，但我收到了几个警告，distributed.worker - WARNING - Memory use is high but worker has no data to store to disk. Perhaps some other process is leaking memory? Process memory: 12.85 GiB -- Worker memory limit: 7.45 GiB，并且数据没有在 GPU 上处理。

例如：

import cupy as cp
import numpy as np
import dask.array as da
from dask_image import ndfilters as dfilters
from dask.distributed import Client
from functools import partial


if __name__ == '__main__':

    client = Client(memory_limit='8GB', processes=False)
    arr = da.from_array(np.zeros((50, 256, 512, 512), dtype=np.uint16), chunks=(1, 64, 256, 256))
    arr = arr.map_blocks(cp.asarray)

    filtering = partial(dfilters.gaussian_filter, sigma=2)

    scattered_data = client.scatter(arr)
    sent = client.submit(filtering, scattered_data)

    filtered = sent.result().compute()
    client.close()

GPU 有 24GB 内存。

提前致谢。

Answer 1

回答具体问题：不，Dask 无法知道或控制任务内部将使用多少内存。从 Dask 的角度来看，这是任意代码，它只是由 python “调用”。在单独的线程中监控总进程内存是最好的工具。

-以前-

不要这样做：

da.from_array(np.zeros((50, 256, 512, 512), dtype=np.uint16), chunks=(1, 64, 256, 256))

您正在实现一个大型数组，将其切碎并运送给工作人员，在使用前需要对其进行反序列化。如果可以，请始终在工作人员中创建数据，在这种简单化的情况下，这相当于

da.zeros((50, 256, 512, 512), dtype=np.uint16), chunks=(1, 64, 256, 256)

或者在 zarr 的情况下使用da.from_zarr。