dask_cuda_local_cluster_gpu_detection_quickstart.py

1from dask_cuda import LocalCUDACluster
2from dask.distributed import Client
3
4# Create a Dask cluster that uses all available GPUs in the local machine
5# Each GPU will be assigned to a single Dask worker
6cluster = LocalCUDACluster()
7
8# Connect a Dask client to the cluster to start submitting tasks
9client = Client(cluster)
10
11# Check the client status to see the number of workers/GPUs detected
12print(client)
13
14# Example: Inspect the resources of the workers
15def get_gpu_info():
16    import pynvml
17    pynvml.nvmlInit()
18    handle = pynvml.nvmlDeviceGetHandleByIndex(0)
19    name = pynvml.nvmlDeviceGetName(handle)
20    return name
21
22# Run a simple function on all GPU workers
23results = client.run(get_gpu_info)
24print(results)