4

I'm setting up a Dask script to be executed on PSC Bridges P100 GPU nodes. These nodes offer 2 GPUs and 32 CPU-cores. I would like to start CPU and GPU-based dask-workers.

The CPU workers will be started:

dask-worker --nprocs 1 --nthreads 1

while the GPU workers as:

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

My workflow consists of a set of CPU only tasks and a set of GPU tasks, that depend on the results from the CPU tasks. Is there a way to bind a GPU task only to the GPU workers?

Furthermore, I would like to make sure that the GPU tasks land on the same compute node as the CPU task they depend on. Can I somehow do that?

underscore_d
  • 6,309
  • 3
  • 38
  • 64
neiron21
  • 71
  • 5

1 Answers1

0

For your sort of problem it makes sense to run dask using the dask.distributed backend (a more sophisticated task scheduler), which provides a functionality called "worker resources".

For every worker it lets you specify virtual worker resources with an associated count, such as "GPU=2". On the client side you can then specify which and how many resources are required for each task. See the docs here.

For making sure a GPU task lands on the same compute node as the task it depends on you could:

  • set the resources accordingly, i.e. splitting up tasks explicitly using resources like "GPU1" and "GPU2"
  • alternatively, bundle the CPU and subsequent GPU task into one task either by manually defining an encompassing function or by using dask graph optimizations as documented here (I'm thinking of "fusing" tasks).
malbert
  • 308
  • 1
  • 7