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 worker
environ={'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: