UCX 集成

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

硬件要求

要使用UCX与NVLink或InfiniBand，相关的GPU必须分别通过NVLink桥接器或NVIDIA Mellanox InfiniBand适配器连接。 NVIDIA提供了NVLink桥接器和InfiniBand适配器的比较图表。

软件需求

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

使用UCX时，每个NVLink和InfiniBand内存缓冲区必须在它们传输的每对唯一进程之间创建映射；这可能相当昂贵，每个映射可能需要数百毫秒。因此，强烈建议使用RAPIDS内存管理器（RMM）来分配一个仅需一次映射操作的内存池，所有后续传输都可以依赖此内存池。内存池还可以防止Dask调度程序反序列化CUDA数据，这会导致崩溃。

警告

Dask-CUDA 必须在集群初始化期间创建工作者 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）。也可以分别启动集群——调度器、工作器和客户端作为不同的进程——只要Dask调度器是用dask scheduler --protocol="ucx"创建的，并且将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。

将sockcm替换为rdmacm在UCX_SOCKADDR_TLS_PRIORITY中，启用远程直接内存访问（RDMA）用于InfiniBand传输。 UCX建议在使用InfiniBand时使用此设置，如果InfiniBand被禁用，则此设置将不起作用。
distributed.rmm.pool-size: – 推荐。

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

注意

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

用法

请参阅启用UCX通信以了解UCX与不同支持的传输方式的使用示例。

在fork资源匮乏的环境中运行

许多高性能网络堆栈不支持用户应用程序在网络底层初始化后调用fork()。症状包括作业随机挂起或崩溃，尤其是在使用大量工作进程时。为了在使用Dask-CUDA的UCX集成时缓解这个问题，通过多进程启动的进程应使用“forkserver”方法启动进程。当使用dask cuda 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 进行子进程调用

注意

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