如何将任务分配给 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 任务位于同一个计算节点上。我能做到吗？

Answer 1

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

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

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

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