如何为GPU和CPU Dask工作人员分配任务？

Question

我正在设置一个Dask脚本，以便在PSC Bridges P100 GPU节点上执行。这些节点提供2个GPU和32个CPU内核。我想开始使用基于CPU和GPU的dask-workers。

将启动CPU工作程序：

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后端（功能更复杂的任务计划程序）运行dask.distributed很有意义，该后端提供了称为“工作者资源”的功能

对于每个工作程序，您都可以指定具有相关计数的虚拟工作程序资源，例如“ GPU = 2”。然后，可以在客户端上指定每个任务需要哪些资源以及需要多少资源。请参阅文档here。

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

• 相应地设置资源，即明确划分任务 使用“ GPU1”和“ GPU2”之类的资源
• 或者，通过手动定义包围功能或通过使用here中所述的模糊图形优化将CPU和后续GPU任务捆绑为一个任务（我正在考虑“融合”任务）。