dask-cuda

Dask-CUDA 是一个库，它扩展了 dask.distributed 中的 LocalCluster，以实现多 GPU 工作负载。

LocalCUDACluster

您可以使用 LocalCUDACluster 在您的本地机器上创建一个包含一个或多个 GPU 的集群。您可以在 LocalCUDACluster 上启动一个 Dask 调度器，以在单个节点上的多个 GPU 上并行化和分发您的 RAPIDS 工作流。

除了支持多 GPU 计算外，LocalCUDACluster 还提供了一个简单的接口来管理集群，例如启动和停止集群、查询节点状态以及监控工作负载分布。

先决条件

在运行这些说明之前，请确保您已在本地环境中安装了 dask 和 dask-cuda 包。

集群设置

实例化 LocalCUDACluster 对象

LocalCUDACluster 类会自动检测系统中的 GPU，因此如果您在具有两个 GPU 的机器上创建它，它将创建一个包含两个 worker 的集群，每个 worker 负责在单独的 GPU 上执行任务。

from dask_cuda import LocalCUDACluster
from dask.distributed import Client

cluster = LocalCUDACluster()

您还可以通过设置 CUDA_VISIBLE_DEVICES 环境变量或将其作为关键字参数来限制您的集群使用特定的 GPU。

cluster = LocalCUDACluster(
    CUDA_VISIBLE_DEVICES="0,1"
)  # Creates one worker for GPUs 0 and 1

连接 Dask 客户端

Dask 调度器协调任务的执行，而 Dask 客户端是面向用户的接口，用于向调度器提交任务并监控其进度。

client = Client(cluster)

测试 RAPIDS

要测试 RAPIDS，请为集群创建一个 distributed 客户端并查询 GPU 模型。

def get_gpu_model():
    import pynvml

    pynvml.nvmlInit()
    return pynvml.nvmlDeviceGetName(pynvml.nvmlDeviceGetHandleByIndex(0))


result = client.submit(get_gpu_model).result()

print(result)
# b'Tesla V100-SXM2-16GB