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:

On this page