dask - 如何将任务分配给 GPU 和 CPU Dask 工作人员？

Question

我正在设置要在 PSC Bridges P100 GPU 节点上执行的 Dask 脚本。这些节点提供 2 个 GPU 和 32 个 CPU 内核。我想启动基于 CPU 和 GPU 的 dask-workers。

CPU worker 将被启动：

dask-worker --nprocs 1 --nthreads 1

而 GPU 工作人员为：

CUDA_VISIBLE_DEVICE=0 dask-worker --nprocs 1 --nthreads 1

我的工作流程由一组仅 CPU 任务和一组 GPU 任务组成，这取决于 CPU 任务的结果。有没有办法将 GPU 任务仅绑定到 GPU 工作人员？

此外，我想确保 GPU 任务与它们所依赖的 CPU 任务位于相同的计算节点上。我能以某种方式做到这一点吗？

标签： dask

Answer 1

dask对于您的问题，使用dask.distributed后端（更复杂的任务调度程序）运行是有意义的，它提供了一种称为“工作资源”的功能。

对于每个工作人员，它允许您指定具有关联计数的虚拟工作人员资源，例如“GPU = 2”。然后，您可以在客户端指定每个任务需要哪些资源以及需要多少资源。请参阅此处的文档。

为了确保 GPU 任务与它所依赖的任务位于相同的计算节点上，您可以：

• 相应地设置资源，即使用“GPU1”和“GPU2”等资源显式拆分任务
• 或者，通过手动定义包含函数或使用此处记录的 dask 图优化（我正在考虑“融合”任务），将 CPU 和后续 GPU 任务捆绑到一个任务中。