UCX 集成

通信可能是分布式系统中的主要瓶颈。Dask-CUDA 通过支持与 UCX 集成来解决此问题，UCX 是一个优化的通信框架，提供高性能网络连接，并支持多种传输方法，包括适用于专用硬件系统的 NVLink 和适用于没有专用硬件的系统的 InfiniBand，以及 TCP。这种集成通过 UCX-Py 实现，UCX-Py 是一个为 UCX 提供 Python 绑定的接口。

硬件要求

要将 UCX 与 NVLink 或 InfiniBand 一起使用，必须分别使用 NVLink bridges 或 NVIDIA Mellanox InfiniBand Adapters 连接相应的 GPU。NVIDIA 提供了 NVLink bridges 和 InfiniBand adapters 的对比图表。

软件要求

UCX 集成需要安装了 UCX 和 UCX-Py 的环境；有关此过程的详细说明，请参阅 UCX-Py 安装。

使用 UCX 时，每个 NVLink 和 InfiniBand 内存缓冲区必须在它们传输的每对唯一进程之间创建映射；这可能非常耗时，每次映射可能长达数百毫秒。因此，强烈建议使用 RAPIDS 内存管理器 (RMM) 分配一个内存池，该内存池仅容易受到单个映射操作的影响，所有后续传输都可以依赖此池。内存池还可以防止 Dask 调度程序反序列化 CUDA 数据，这会导致崩溃。

警告

Dask-CUDA 必须在集群初始化期间创建 worker CUDA 上下文，正确安排该任务对于正确的 UCX 配置至关重要。如果在集群初始化时该进程已存在 CUDA 上下文，则可能会发生意外行为。为了避免这种情况，建议在执行会导致创建 CUDA 上下文的操作之前初始化任何启用 UCX 的集群。根据库的不同，即使是导入也可能强制创建 CUDA 上下文。

对于某些 RAPIDS 库（例如 cuDF），在运行时设置 RAPIDS_NO_INITIALIZE=1 将延迟或禁用其 CUDA 上下文创建，从而提高与启用 UCX 的集群的兼容性并防止运行时警告。

配置

自动配置

从 Dask-CUDA 22.02 版本开始，并且假设 UCX >= 1.11.1，现在指定 UCX 传输方式是可选的。

现在可以使用 LocalCUDACluster(protocol="ucx") 启动本地集群，这意味着自动选择 UCX 传输方式 (UCX_TLS=all)。分开启动集群——将调度程序、worker 和客户端作为不同的进程——也是可能的，只要使用 dask scheduler --protocol="ucx" 创建 Dask 调度程序，并且将 dask cuda worker 连接到调度程序将意味着自动选择 UCX 传输方式，但这需要 Dask 调度程序和客户端使用 DASK_DISTRIBUTED__COMM__UCX__CREATE_CUDA_CONTEXT=True 启动。有关 UCX 与自动配置一起使用的更多详细示例，请参阅启用 UCX 通信。

仍然可以手动配置传输方式，请参阅以下小节。

手动配置

除了在系统上安装 UCX 和 UCX-Py 之外，对于手动配置，必须在 Dask 配置中指定多个选项以启用集成。通常，这些选项将影响 UCX_TLS 和 UCX_SOCKADDR_TLS_PRIORITY，这些是 UCX 用于决定使用哪种传输方法以及优先使用哪种方法的环境变量。但是，有些会影响相关的库，例如 RMM

distributed.comm.ucx.cuda_copy: true – 必需。

将 cuda_copy 添加到 UCX_TLS，启用通过 UCX 进行 CUDA 传输。
distributed.comm.ucx.tcp: true – 必需。

将 tcp 添加到 UCX_TLS，启用通过 UCX 进行 TCP 传输；这对于非常小的传输是必需的，因为 NVLink 和 InfiniBand 处理小传输效率不高。
distributed.comm.ucx.nvlink: true – NVLink 必需。

将 cuda_ipc 添加到 UCX_TLS，启用通过 UCX 进行 NVLink 传输；仅影响节点内部通信。
distributed.comm.ucx.infiniband: true – InfiniBand 必需。

将 rc 添加到 UCX_TLS，启用通过 UCX 进行 InfiniBand 传输。

为了在 UCX 1.11 及更高版本中获得最佳性能，建议同时设置环境变量 UCX_MAX_RNDV_RAILS=1 和 UCX_MEMTYPE_REG_WHOLE_ALLOC_TYPES=cuda，有关这些变量的更多详细信息，请参阅此处和此处的文档。
distributed.comm.ucx.rdmacm: true – InfiniBand 推荐。

在 UCX_SOCKADDR_TLS_PRIORITY 中用 rdmacm 替换 sockcm，为 InfiniBand 传输启用远程直接内存访问 (RDMA)。UCX 推荐将其用于 InfiniBand，如果 InfiniBand 被禁用，则此选项无效。
distributed.rmm.pool-size: <str|int> – 推荐。

为进程分配指定大小的 RMM 池；大小可以提供整数字节数或易读格式，例如 "4GB"。建议将池大小设置为至少为进程使用的最小内存量；如果可能，可以将所有 GPU 内存映射到单个池，以供进程的生命周期使用。

注意

这些选项可以与主线 Dask.distributed 一起使用。但是，某些功能是 Dask-CUDA 独有的，例如 InfiniBand 接口的自动检测。有关使用 Dask-CUDA 的好处的更多详细信息，请参阅Dask-CUDA – 动机。

使用方法

有关 UCX 与不同支持的传输方式一起使用的示例，请参阅启用 UCX 通信。

在 fork 资源受限环境中运行

许多高性能网络栈不支持用户应用程序在网络底层初始化后调用 fork()。症状包括作业随机挂起或崩溃，尤其是在使用大量 worker 时。为了在使用 Dask-CUDA 的 UCX 集成时缓解此问题，通过多进程启动的进程应使用 “forkserver” 方法启动进程。使用 dask cuda worker 启动 worker 时，可以通过传递参数 --multiprocessing-method forkserver 来实现。在用户代码中，可以使用 dask 中的 distributed.worker.multiprocessing-method 配置键来控制此方法。此外，必须注意手动确保 forkserver 在启动任何作业之前正在运行。因此，运行脚本应执行类似于以下操作：

import dask

if __name__ == "__main__":
    import multiprocessing.forkserver as f
    f.ensure_running()
    with dask.config.set(
        {"distributed.worker.multiprocessing-method": "forkserver"}
    ):
        run_analysis(...)

注意

除此之外，目前还必须在环境中设置 PTXCOMPILER_CHECK_NUMBA_CODEGEN_PATCH_NEEDED=0 以避免 ptxcompiler 中的子进程调用

注意

要确认没有发生错误的 fork 调用，请使用 UCX_IB_FORK_INIT=n 启动作业。如果应用程序调用 fork()，UCX 将生成警告 UCX WARN IB: ibv_fork_init() was disabled or failed, yet a fork() has been issued.。

故障排除

超时设置

根据集群大小和 GPU 架构的不同，Dask worker 之间建立端点时可能会发生超时。对于这些情况，可以通过 distributed.comm.ucx.connect-timeout 配置或相应的 DASK_DISTRIBUTED__COMM__UCX__CONNECT_TIMEOUT 环境变量来增加默认超时时间。该值表示超时时间，单位为秒。

请注意，超时设置旨在防止 worker 在出现问题时无限期挂起，因此将超时时间设置得过高可能会导致 worker 看起来冻结。因此，务必谨慎增加此值，并将其保持在合理的短时间内。截至目前，尚未观察到将此值增加到 60 秒仍不足够的情况。