Supporting scheduling on specific GPU

[aadamson]

I'm trying to support the following scenario: We have multiple machines with multiple NVIDIA GPUs. On each machine, we run one worker per GPU. When a task gets scheduled on a worker, we want to make sure that any GPU routine invoked by the task runs on a GPU that is not in use by any other worker. Libraries such as TensorFlow and pytorch will schedule GPU work on the visible GPU with the lowest bus id, so by default, tasks running in parallel will both try to schedule on the same GPU if they have the same visible GPU set. CUDA supports masking available GPUs to processes via setting the CUDA_VISIBLE_DEVICES environment variable. I want to ensure that if two tasks are running on workers on the same machine, the GPUs they can see (and hence the GPUs they can schedule on) do not overlap.

I can think of two possible approaches here:

• Assigning each worker a GPU to use on startup by running dask-worker with CUDA_VISIBLE_DEVICES set such that each worker has a unique GPU, and --resources "GPU=1"
• Add a scheduler plugin that will assign tasks to available GPUs on the host of the worker they were scheduled on and add a preamble to the task's target that will set the environment variables properly.

I'm new to Dask, so I'm not sure if either of these is appropriate. Is there prior art on handling situations like this?

[mrocklin]

Assigning each worker a GPU to use on startup by running dask-worker with CUDA_VISIBLE_DEVICES set such that each worker has a unique GPU, and --resources "GPU=1"

I believe I've seen this done before. This video is a bit old, but may be helpful: https://youtu.be/jLQ2wAwapRg?t=894

[aadamson]

Thanks Matt, great memory :D I'll take a crack at using that.

[VMois]

@aadamson Have you tried to set CUDA_VISIBLE_DEVICES for each worker? Is it working good?

[gar1t]

This seems like a mainstream use case. Was there any resolution to this? Is there a recommended approach now for scheduling workers on a specific GPU?

[quasiben]

Yes: https://distributed.dask.org/en/latest/resources.html?#example

If you run into problems or have new thoughts I would recommend opening up a new issue

[gar1t]

My understanding of N in a resource GPU=N is that it's the number of GPUs (a count), not a specific GPU. The issue here as I read it is that the OP wants to schedule a worker on a specific GPU. If you schedule two workers with CUDA_VISIBLE_DEVICES=2 and a Dask resource of GPU=1, does Dask limit the pool of GPUs to a single GPU with ID 2 in that case?

[gar1t]

Having read through the source code I believe my understanding here is correct. The resource is a counter that limits what is scheduled per worker. As this issue describes the problem accurately (the title and original description are crystal clear) I don't see the point in creating a new issue. The question stands and is not address here. I suggest simply re-opening it or otherwise resolving it (e.g. with an answer or a won't fix, etc.)

[gar1t]

Okay, just to cap this off... This can be implemented by using GPU device specific resource names with value 1. E.g. when starting a worker. E.g. for a two GPU system:

$ dask-worker scheduler:8786 --resources "gpu0=1"
$ dask-worker scheduler:8786 --resources "gpu1=1"

When submitting a task for the first GPU:

>>> client.submit(task_configured_to_use_gpu0, resource={"gpu0": 1})
>>> client.submit(task_configured_to_use_gpu1, resource={"gpu1": 1})

In this example, the respective tasks must ensure that they run on the target GPU (e.g. by setting CUDA_VISIBLE_DEVICES for an external process, TensorFlow device context manager, etc.)