UCX 扩展#

描述#

UCX 扩展利用统一通信 X (UCX) 库在 GXF 框架中分离图。此扩展有助于跨多个主机分布图，从而能够利用分布式 GPU 资源。UCX 是一个开源库，以其加速跨高性能网络的数据传输的能力而闻名。它可以利用 GPUDirect RDMA 技术来优化网络延迟并最大化分布式 GPU 流量吞吐量。因此，此扩展的用户可以利用跨不同主机的多个 GPU 的组合处理能力。这可以显著提高工作流程的速度和效率。有关 UCX 的更多详细信息，请访问 https://openucx.org。

例如

下图说明了一个分离的图，该图由两个张量生成器和一个张量比较器组成。此张量比较器评估这些张量生成器产生的输出。UcxExtension 提供了在不同主机上执行每个实体的能力。

为此，每个使用 UCX 扩展的图都需要一个 UcxContext 组件。此组件托管 UCP 上下文并负责所有连接、管理数据，并确保所有操作在反初始化时正确关闭。在设置图时，请将实体的标准发送器和接收器替换为 UcxTransmitter 和 UcxReceiver 组件。请务必配置所有参数，包括 IP、端口和其他参数，以正确建立连接。

目前，UCX 支持发送相同类型的内存（主机或设备）的消息。这是 UCX 的限制，而不是扩展的限制。

UUID: 525f8a1a-dfb5-426b-8ddb-00c3ac839994
版本: 0.8.0
作者: NVIDIA
许可证: LICENSE

要求#

NVIDIA ConnectX6-DX NIC 或更高版本。

有关安装和配置 NIC 的更多信息，请参阅： https://docs.nvda.net.cn/networking/display/ConnectX6VPI/Introduction
Mellanox Open Fabrics Enterprise Distribution (MLNX_OFED) - 5.5 或更高版本，请参阅 https://network.nvidia.com/products/infiniband-drivers/linux/mlnx_ofed/
有关安装说明，请参阅 https://docs.nvda.net.cn/networking/display/MLNXOFEDv551032/Installing+MLNX_OFED
如果在容器中安装 Mellanox OFED
- 确保通过将 --all 标志传递给 mlnxofedinstall 脚本，在主机操作系统中安装内核驱动程序。
- 在容器中，您只能使用 --user-space-only 标志传递给 mlnxofedinstall 脚本来安装用户空间库。
UCX - 1.13 或更高版本 - 需要使用 CUDA 支持进行编译，或直接从 git 存储库中使用启用 CUDA 的 UCX 包，请参阅 openucx/ucx
有关安装说明，请按照此处的发布版本构建说明进行操作： openucx/ucx。

请注意，UCX 库应使用 CUDA 编译，如下所示:: .. code-block:: bash

$ ./contrib/configure-release –prefix=/install/path –enable-examples –with-java=no –with-cuda=/path/to/cuda –enable-mt

组件#

UcxContext#

UcxContext 在 GXF UCX 扩展中至关重要。它负责初始化 UCX 上下文、运行侦听器以及管理 UcxReceiver 的连接请求和数据接收。UcxContext 还设置 UcxTransmitter 连接和资源。所有连接（包括 UcxReceiver 和 UcxTransmitter 的连接）都在 UcxContext 中管理。在图完成后，UcxContext 负责关闭所有连接并释放所有资源。

组件 ID: 755d20a5-d794-467d-a86c-290eb2c32052
基类型: nvidia::gxf::NetworkContext
定义于: extensions/ucx/ucx_context.hpp

参数#

序列化器

组件使用的实体序列化器。应使用 UcxComponentSerializer 类型。

标志: GXF_PARAMETER_FLAGS_NONE
类型: GXF_PARAMETER_TYPE_HANDLE
句柄类型: nvidia::gxf::EntitySerializer

重新连接

如果在运行期间连接关闭，则尝试重新连接。对于 UcxReceiver，它将等待新的连接请求以建立新连接。对于 UcxTransmitter，它将向服务器发送新的连接请求以建立新连接。

标志: GXF_PARAMETER_FLAGS_NONE
类型: GXF_PARAMETER_TYPE_BOOL
默认值: true

可选 GPU 设备资源

GPU 设备的可选资源。

标志: GXF_PARAMETER_FLAGS_NONE
类型: GXF_PARAMETER_TYPE_HANDLE
句柄类型: nvidia::gxf::GPUDevice

UcxTransmitter#

GXF UCX 扩展的发送器组件。此组件用作实体的发送器。在初始化阶段，它将发送连接请求以建立连接。当网络路由器执行 SyncOutbox 函数时，它会调用 UcxTransmitter 的 sync_io 方法。此方法反过来利用 UCX Active Message Rendezvous 协议传输消息。

组件 ID: 58165d03-78b7-4696-b200-71621f90aee7
基类型: nvidia::gxf::Transmitter
定义于: extensions/ucx/ucx_transmitter.hpp

参数#

容量

发送器的队列容量。

标志: GXF_PARAMETER_FLAGS_NONE
类型: GXF_PARAMETER_TYPE_UINT64

策略

队列的数据处理策略。有效值

0: 弹出 1: 拒绝 2: 故障

标志: GXF_PARAMETER_FLAGS_NONE
类型: GXF_PARAMETER_TYPE_UINT64

接收器地址

要连接的接收器地址。

标志: GXF_PARAMETER_FLAGS_NONE
类型: GXF_PARAMETER_TYPE_STRING

端口

接收器的端口。

标志: GXF_PARAMETER_FLAGS_NONE
类型: GXF_PARAMETER_TYPE_INT32

缓冲区

用于保存序列化数据的序列化缓冲区。

标志: GXF_PARAMETER_FLAGS_NONE
类型: GXF_PARAMETER_TYPE_HANDLE
句柄类型: Handle<UcxSerializationBuffer>

最大连接重试次数

建立连接的最大重试次数。

标志: GXF_PARAMETER_FLAGS_NONE
类型: GXF_PARAMETER_TYPE_INT32

gpu_device

可选 GPU 设备资源。

标志: GXF_PARAMETER_FLAGS_NONE
类型: GXF_PARAMETER_TYPE_HANDLE
句柄类型: Handle<GPUDevice>

UcxReceiver#

在 GXF UCX 扩展中接收数据。此组件替换实体的接收器。当实体向此接收器发送消息时，UCXContext 接收消息头，提示路由器执行 SyncInbox 函数。SyncInbox 函数随后触发 UcxReceiver 的 sync_io 方法。此方法利用 UCX Active Message Rendezvous 协议接收消息的数据内容。

组件 ID: e961132b-45d5-48b8-ac5d-2bb1a4a42279
基类型: nvidia::gxf::Receiver
定义于: extensions/ucx/ucx_receiver.hpp

参数#

容量

接收器的队列容量。

标志: GXF_PARAMETER_FLAGS_NONE
类型: GXF_PARAMETER_TYPE_UINT64
默认值: 10

策略

队列的数据处理策略。0: 弹出, 1: 拒绝, 2: 故障

标志: GXF_PARAMETER_FLAGS_NONE
类型: GXF_PARAMETER_TYPE_UINT64
默认值: 2

地址

用于接收数据的侦听器地址。

标志: GXF_PARAMETER_FLAGS_NONE
类型: GXF_PARAMETER_TYPE_STRING
默认值: “0.0.0.0”

端口

用于接收数据的侦听器的端口。

标志: GXF_PARAMETER_FLAGS_NONE
类型: GXF_PARAMETER_TYPE_INT32
默认值: 13337

缓冲区

用于保存序列化数据的序列化缓冲区。

标志: GXF_PARAMETER_FLAGS_NONE
类型: GXF_PARAMETER_TYPE_HANDLE
句柄类型: UcxSerializationBuffer

可选 GPU 设备资源

GPU 设备的可选资源。

标志: GXF_PARAMETER_FLAGS_NONE
类型: GXF_PARAMETER_TYPE_HANDLE
句柄类型: nvidia::gxf::GPUDevice

UcxComponentSerializer#

GXF UCX 扩展中组件的序列化器。目前支持序列化时间戳、张量、视频缓冲区、音频缓冲区和整数组件。适用于在具有相同字节序的设备之间共享数据。

组件 ID: 64994305-4260-4f5c-ac5f-69da6dd6cfa5
基类型: nvidia::gxf::ComponentSerializer
定义于: extensions/ucx/ucx_component_serializer.hpp

参数#

分配器

张量组件的内存分配器。

标志: GXF_PARAMETER_FLAGS_NONE
类型: GXF_PARAMETER_TYPE_HANDLE
句柄类型: nvidia::gxf::Allocator

UcxEntitySerializer#

GXF UCX 扩展中实体的序列化器。

组件 ID: 14997aa4-4a01-4cd4-86ab-687f85a13f10
基类型: nvidia::gxf::EntitySerializer
定义于: extensions/ucx/ucx_entity_serializer.hpp

参数#

组件序列化器

用于序列化和反序列化组件的序列化器列表。

标志: GXF_PARAMETER_FLAGS_NONE
类型: GXF_PARAMETER_TYPE_HANDLE
句柄类型: FixedVector<nvidia::gxf::Handle<nvidia::gxf::ComponentSerializer>, kMaxTempComponents>

详细警告

是否打印详细警告。

标志: GXF_PARAMETER_FLAGS_NONE
类型: GXF_PARAMETER_TYPE_BOOL
默认值: true

UcxSerializationBuffer#

GXF UCX 扩展的序列化缓冲区。

组件 ID: 1d9fcaf7-1db1-4992-93ec-714979f7d78d 基类型: nvidia::gxf::Endpoint 定义于: extensions/ucx/ucx_serialization_buffer.hpp

参数#

分配器

张量组件的内存分配器。

标志: GXF_PARAMETER_FLAGS_NONE
类型: GXF_PARAMETER_TYPE_HANDLE
句柄类型: nvidia::gxf::Handle<nvidia::gxf::Allocator>

缓冲区大小

缓冲区大小（字节），默认为 4kB。

标志: GXF_PARAMETER_FLAGS_NONE
类型: GXF_PARAMETER_TYPE_SIZE
默认值: 4096 (4kB)

示例#

本节提供了一个在简单图中使用 UCX 扩展的示例。此图包含两个子图，通过 UCX 扩展互连。服务器端和客户端的配置详细信息封装在各自的 YAML 文件中，这些文件在下面共享以供参考。

服务器端 - test_ping_rx.yaml 文件

name: rx
components:
- name: allocator
  type: nvidia::gxf::test::MockAllocator
- name: serialization_buffer
  type: nvidia::gxf::UcxSerializationBuffer
  parameters:
    allocator: allocator
- name: signal
  type: nvidia::gxf::UcxReceiver
  parameters:
    address: 5.5.5.5
    port: 13337
    buffer: serialization_buffer
- type: nvidia::gxf::MessageAvailableSchedulingTerm
  parameters:
    receiver: signal
    min_size: 1
- type: nvidia::gxf::PingRx
  parameters:
    signal: signal
- type: nvidia::gxf::test::StepCount
  parameters:
    expected_count: 10
- type: nvidia::gxf::CountSchedulingTerm
  parameters:
    count: 10
---
name: ucx
components:
- name: allocator
  type: nvidia::gxf::test::MockAllocator
- name: component_serializer
  type: nvidia::gxf::UcxComponentSerializer
  parameters:
    allocator: allocator
- name: entity_serializer
  type: nvidia::gxf::UcxEntitySerializer
  parameters:
    component_serializers: [ component_serializer ]
- name: ucx_context
  type: nvidia::gxf::UcxContext
  parameters:
    serializer: entity_serializer
---
name: scheduler
components:
- name: clock
  type: nvidia::gxf::RealtimeClock
- type: nvidia::gxf::GreedyScheduler
  parameters:
    max_duration_ms: 1000000
    clock: clock
    stop_on_deadlock: False
---
name: gpu_resource_entity_0
components:
- type: nvidia::gxf::GPUDevice
  name: gpu_resource_0
  parameters:
    dev_id: 0
---
EntityGroups:
- name: entity_group_0
  target:
  - "rx"
  - "ucx"
  - "gpu_resource_entity_0"

客户端 - test_ping_tx.yaml 文件

name: tx
components:
- name: allocator
  type: nvidia::gxf::test::MockAllocator
- name: serialization_buffer
  type: nvidia::gxf::UcxSerializationBuffer
  parameters:
    allocator: allocator
- name: signal
  type: nvidia::gxf::UcxTransmitter
  parameters:
    receiver_address: 5.5.5.5
    port: 13337
    buffer: serialization_buffer
- type: nvidia::gxf::PingTx
  parameters:
    signal: signal
- type: nvidia::gxf::CountSchedulingTerm
  parameters:
    count: 10
- type: nvidia::gxf::test::StepCount
  parameters:
    expected_count: 10
---
name: ucx
components:
- name: allocator
  type: nvidia::gxf::test::MockAllocator
- name: component_serializer
  type: nvidia::gxf::UcxComponentSerializer
  parameters:
    allocator: allocator
- name: entity_serializer
  type: nvidia::gxf::UcxEntitySerializer
  parameters:
    component_serializers: [ component_serializer ]
- name: ucx_context
  type: nvidia::gxf::UcxContext
  parameters:
    serializer: entity_serializer
---
name: scheduler
components:
- name: clock
  type: nvidia::gxf::RealtimeClock
- type: nvidia::gxf::GreedyScheduler
  parameters:
    stop_on_deadlock: false
    max_duration_ms: 1000000
    clock: clock
---
name: gpu_resource_entity_0
components:
- type: nvidia::gxf::GPUDevice
  name: gpu_resource_0
  parameters:
    dev_id: 0
---
EntityGroups:
- name: entity_group_0
  target:
  - "tx"
  - "ucx"
  - "gpu_resource_entity_0"