启用 UCX 通信

使用 UCX 通信的 CUDA 集群可以通过 LocalCUDACluster 自动启动，或通过 dask cuda worker CLI 工具手动启动。无论哪种情况，都必须为工作进程集群创建使用相同 Dask UCX 配置的 dask.distributed.Client；有关所有可用选项的详细信息，请参阅UCX 集成 – 配置。

使用自动配置的 LocalCUDACluster

自动配置在 Dask-CUDA 22.02 中引入，需要 UCX >= 1.11.1。这允许用户仅指定 UCX 协议，让 UCX 决定使用哪些传输方式。

要将客户端连接到使用自动配置的 UCX 和 RMM 池的集群

from dask.distributed import Client
from dask_cuda import LocalCUDACluster

cluster = LocalCUDACluster(
    protocol="ucx",
    interface="ib0",
    rmm_pool_size="1GB"
)
client = Client(cluster)

注意

上面特意指定 interface="ib0" 以确保在支持 InfiniBand 的系统中使用了 RDMACM。在不支持 InfiniBand 或不需要 RDMACM 的系统上，可以省略 interface 参数，或指定监听其他接口。

使用手动配置的 LocalCUDACluster

当使用手动配置的 LocalCUDACluster 和 UCX 通信时，所有必需的 UCX 配置都通过构造函数提供的参数处理；有关这些参数的完整列表，请参阅API – Cluster。要将客户端连接到支持所有传输方式和 RMM 池的集群

from dask.distributed import Client
from dask_cuda import LocalCUDACluster

cluster = LocalCUDACluster(
    protocol="ucx",
    interface="ib0",
    enable_tcp_over_ucx=True,
    enable_nvlink=True,
    enable_infiniband=True,
    enable_rdmacm=True,
    rmm_pool_size="1GB"
)
client = Client(cluster)

使用自动配置的 dask cuda worker

当使用自动配置的 dask cuda worker 和 UCX 通信时，调度器、工作进程和客户端都必须手动启动，但无需显式指定任何 UCX 传输方式。这仅在 Dask-CUDA 22.02 及更新版本中支持，并且需要 UCX >= 1.11.1。

调度器

对于自动 UCX 配置，我们必须确保在 UCX 初始化之前，调度器上已创建 CUDA 上下文。这可以通过在创建调度器时指定环境变量 DASK_DISTRIBUTED__COMM__UCX__CREATE_CUDA_CONTEXT=True 来实现。

要启动使用 UCX 进行自动配置和 1GB RMM 池的 Dask 调度器

$ DASK_DISTRIBUTED__COMM__UCX__CREATE_CUDA_CONTEXT=True \
> DASK_DISTRIBUTED__RMM__POOL_SIZE=1GB \
> dask scheduler --protocol ucx --interface ib0

注意

上面特意指定 interface="ib0" 以确保在支持 InfiniBand 的系统中使用了 RDMACM。在不支持 InfiniBand 或不需要 RDMACM 的系统上，可以省略 interface 参数，或指定监听其他接口。

上面我们指定了 UCX_MEMTYPE_REG_WHOLE_ALLOC_TYPES=cuda 以在 InfiniBand 环境中获得最佳性能，详细信息请参阅此处。如果不使用 InfiniBand，可以省略该选项。在 UCX 1.12 及更高版本中，该选项是默认的，即使在使用 InfiniBand 时也可以省略。

工作进程

要启动使用 UCX 进行自动配置和每个 GPU 14GB RMM 池的工作进程

$ UCX_MEMTYPE_REG_WHOLE_ALLOC_TYPES=cuda
> dask cuda worker ucx://<scheduler_address>:8786 \
> --rmm-pool-size="14GB" \
> --interface="ib0"

注意

类似于调度器的设置，上面特意指定了 interface="ib0" 以确保在支持 InfiniBand 的系统中使用了 RDMACM。在不支持 InfiniBand 或不需要 RDMACM 的系统上，可以省略 interface 参数，或指定监听其他接口。

上面我们指定了 UCX_MEMTYPE_REG_WHOLE_ALLOC_TYPES=cuda 以在 InfiniBand 环境中获得最佳性能，详细信息请参阅此处。如果不使用 InfiniBand，可以省略该选项。在 UCX 1.12 及更高版本中，该选项是默认的，即使在使用 InfiniBand 时也可以省略。

客户端

要将客户端连接到我们启动的使用 UCX 进行自动配置的集群

import os

os.environ["UCX_MEMTYPE_REG_WHOLE_ALLOC_TYPES"] = "cuda"

import dask
from dask.distributed import Client

with dask.config.set({"distributed.comm.ucx.create_cuda_context": True}):
    client = Client("ucx://<scheduler_address>:8786")

或者，可以省略上面示例中的 with dask.config.set 语句，而改用环境变量 DASK_DISTRIBUTED__COMM__UCX__CREATE_CUDA_CONTEXT=True

import os

os.environ["UCX_MEMTYPE_REG_WHOLE_ALLOC_TYPES"] = "cuda"
os.environ["DASK_DISTRIBUTED__COMM__UCX__CREATE_CUDA_CONTEXT"] = "True"

from dask.distributed import Client

client = Client("ucx://<scheduler_address>:8786")

注意

上面我们指定了 UCX_MEMTYPE_REG_WHOLE_ALLOC_TYPES=cuda 以在 InfiniBand 环境中获得最佳性能，详细信息请参阅此处。如果不使用 InfiniBand，可以省略该选项。在 UCX 1.12 及更高版本中，该选项是默认的，即使在使用 InfiniBand 时也可以省略。

使用手动配置的 dask cuda worker

当使用手动配置的 dask cuda worker 和 UCX 通信时，调度器、工作进程和客户端都必须手动启动，并且每个都使用相同的 UCX 配置。

调度器

UCX 配置选项需要作为环境变量为 dask scheduler 指定；有关环境变量和选项之间映射的更多详细信息，请参阅Dask 配置 – 环境变量。

要启动使用 UCX 并支持所有传输方式和千兆字节 RMM 池的 Dask 调度器

$ DASK_DISTRIBUTED__COMM__UCX__CUDA_COPY=True \
> DASK_DISTRIBUTED__COMM__UCX__TCP=True \
> DASK_DISTRIBUTED__COMM__UCX__NVLINK=True \
> DASK_DISTRIBUTED__COMM__UCX__INFINIBAND=True \
> DASK_DISTRIBUTED__COMM__UCX__RDMACM=True \
> DASK_DISTRIBUTED__RMM__POOL_SIZE=1GB \
> dask scheduler --protocol ucx --interface ib0

我们通过 --protocol 选项通知调度器我们将使用 UCX，并通过 --interface 选项通知它我们将使用 InfiniBand。

工作进程

所有 UCX 配置选项在 dask cuda worker 中都有类似的选项；有关这些选项的完整列表，请参阅API – Worker。要启动支持所有传输方式和 RMM 池的集群

$ dask cuda worker ucx://<scheduler_address>:8786 \
> --enable-tcp-over-ucx \
> --enable-nvlink \
> --enable-infiniband \
> --enable-rdmacm \
> --rmm-pool-size="1GB"

客户端

可以通过使用 dask_cuda.initialize 来配置客户端使用 UCX，这是一个实用程序，它接受与 LocalCUDACluster 相同的 UCX 配置参数，并将其添加到创建客户端时使用的当前 Dask 配置中；有关参数的完整列表，请参阅API – Client initialization。要连接客户端到我们创建的集群

from dask.distributed import Client
from dask_cuda.initialize import initialize

initialize(
    enable_tcp_over_ucx=True,
    enable_nvlink=True,
    enable_infiniband=True,
    enable_rdmacm=True,
)
client = Client("ucx://<scheduler_address>:8786")