RAPIDS#
This is a short example for getting started with GPUs using Coiled and RAPIDS.
You can use the publicly available RAPIDS image with the container
argument and request GPU-enabled machines with the worker_gpu
argument:
import coiled
cluster = coiled.Cluster(
container="nvcr.io/nvidia/rapidsai/rapidsai-core:23.04-cuda11.8-runtime-ubuntu22.04-py3.10",
n_workers=4,
worker_gpu=1, # single T4 per worker
worker_class="dask_cuda.CUDAWorker", # recommended
environ={"DISABLE_JUPYTER": "true"}, # needed for "stable" RAPIDS image
worker_options={"nthreads": 1}, # launch one task per worker. Avoids oversaturating the GPU
)
You’ll notice a couple arguments we recommend, but are not required:
worker_class='dask_cuda.CUDAWorker'
can be helpful if you’re using multiple GPUs per workerenviron={'DISABLE_JUPYTER': 'true'}
is only needed if you are using the stable RAPIDS image and disables the default Jupyter server from starting
Verification#
You can then verify this cluster is working as expected.
from dask.distributed import Client
def test_gpu():
import numpy as np
import cupy as cp
x = cp.arange(6).reshape(2, 3).astype("f")
result = x.sum()
return cp.asnumpy(result), str(result.device)
client = Client(cluster)
f = client.submit(test_gpu)
f.result()
This should return (array(15., dtype=float32), '<CUDA Device 0>')
. You can also verify workers are using GPUs with cluster.scheduler_info["workers"]
.
cluster.scheduler_info["workers"]
And close your cluster when you’re done.
client.shutdown()
Next steps#
For more examples of what you can do with a GPU cluster, see: