Azure VM 集群（通过 Dask）#

使用 Dask Cloud Provider 创建集群#

在 Azure 上设置多节点、多 GPU 集群最简单的方法是使用 Dask Cloud Provider。

1. 安装 Dask Cloud Provider#

Dask Cloud Provider 可以通过 conda 或 pip 安装。Azure 特定的功能需要通过 [azure] pip 扩展进行安装。

$ pip install dask-cloudprovider[azure]

2. 配置 Azure 资源#

根据 Dask Cloud Provider 的说明，设置您的 Azure 资源组、虚拟网络和安全组。

3. 创建集群#

在 Python 终端中，可以使用 dask_cloudprovider 包创建集群。下面的示例在 westus2 区域使用 Standard_NC12s_v3 VM 创建了一个包含 2 个工作节点的集群。VM 应至少有 100GB 的磁盘空间，以容纳 RAPIDS 容器镜像和相关依赖项。

from dask_cloudprovider.azure import AzureVMCluster

resource_group = "<RESOURCE_GROUP>"
vnet = "<VNET>"
security_group = "<SECURITY_GROUP>"
subscription_id = "<SUBSCRIPTION_ID>"
cluster = AzureVMCluster(
    resource_group=resource_group,
    vnet=vnet,
    security_group=security_group,
    subscription_id=subscription_id,
    location="westus2",
    vm_size="Standard_NC12s_v3",
    public_ingress=True,
    disk_size=100,
    n_workers=2,
    worker_class="dask_cuda.CUDAWorker",
    docker_image="nvcr.io/nvidia/rapidsai/base:25.04-cuda12.8-py3.12",
    docker_args="-p 8787:8787 -p 8786:8786",
)

4. 测试 RAPIDS#

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

from dask.distributed import Client

client = Client(cluster)

def get_gpu_model():
    import pynvml

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

client.submit(get_gpu_model).result()

Out[5]: b'Tesla V100-PCIE-16GB'

5. 清理#

使用完集群后，请确保 cluster 和 client 已关闭

client.close()
cluster.close()