We are currently using Dask Gateway with CPU-only workers. However, down the road when deep learning becomes more widely adopted, we want to transition into adding GPU support for the clusters created through Dask Gateway.
I've checked the Dask Gateway documentation, and there isn't so much in the way of detailed instruction on how to set this up and what parts of the helm chart/config we need to change to enable this functionality.
What I'm thinking is to first add a GPU to the GKE cluster on GCP, then use a RAPIDS dockerfile for the dask workers that utilizes this GPU? Is that all the set-up needed for Dask Gateway?
Would appreciate if someone could point me in the right direction.
To run a Dask cluster on Kubernetes capable of GPU compute you need the following:
resources.limits.nvidia.com/gpu: 1
dask-cuda-worker
command from thedask_cuda
package (which is included in the RAPIDS images).Note: For Dask Gateway your container image also needs the
dask-gateway
package to be installed. We can configure this to be installed at runtime but it's probably best to create a custom image with this package installed.Therefore here is a minimal Dask Gateway config which will get you a GPU cluster.
We can test things work by launching Dask gateway, creating a Dask cluster and running some GPU specific work. Here is an example where we get the NVIDIA driver version from each worker.