如何利用 GPU 将大型 dask 数组（numpy.ndarray）写入 Zarr 文件？

Question

我正在尝试使用 dask.如果我的 dask 数组被命名为 dask_data，那么一个简单的 dask_data.to_zarr("my_zarr.zarr") 就可以了。但据我了解，这是一个同步的、受 CPU 限制的进程。

我想做的是使用并行处理分配给 Quadro GV100 GPU 的大部分工作。我尝试通过dask_data_cupy = dask_data.map_blocks(cupy.asarray) 将 numpy.ndarray 转换为 cupy.ndarray 并将其写入 zarr 文件，但我收到：

ValueError: object __array__ method not producing an array（坦率地说，我没有看到性能提升要么）。

如何使用 GPU 将 dask 数组并行写入 zarr 文件？

谢谢！

Answer 1

但据我了解，这是一个同步的、受 CPU 限制的进程。

这可能不是真的，您的瓶颈可能是存储设备。在任何情况下，每个块都被写入一个单独的文件，并跨线程和/或进程并行写入（取决于您的设置）。这就是 zarr 设计的重点，即应用程序可以独立地与每个块进行交互。

如果您选择使用各种压缩编码，您可能会受到 CPU 限制；然而，这些并不一定适合 GPU 操作。

简而言之，除非您的数据已经在 GPU 上生成，否则如果在将其写入文件之前将其传输到 GPU 进行处理是值得的，我会感到惊讶。如果有一个函数可以直接读取/写入 zarr 的 cupy 数组，并且您也在 GPU 上进行处理，那将是不同的 - 但我不相信有。

Answer 2

我认为您需要在调用to_zarr 之前添加dask_data.map_blocks(cupy.asnumpy)。

CuPy 试图确保用户打算进行设备到主机的传输（因为这些可能很昂贵）。因此，当在 CuPy 数组上调用 numpy.asarray 时故意引发（在此写入期间会发生）。