CUDA 零拷贝 IPC

首先假设您在 cuda 上有一个 pytorch tensor，并且您希望在节点之间共享它。

好消息是，您可以使用 CUDA 零拷贝 IPC 在不复制数据的情况下完成此操作。

安装

使用此特性，确保满足以下要求：

# 安装具有 Gpu 支持的 pyarrow
conda install pyarrow "arrow-cpp-proc=*=cuda" -c conda-forge

## 使用如下测试安装
python -c "import pyarrow.cuda"

# 安装从 arrow 至 torchl 转译的 numba
pip install numba

## 使用如下测试安装
python -c "import numba.cuda"

# 如果 torch 不存在请先安装
pip install torch

## 使用如下测试安装
python -c "import torch; assert torch.cuda.is_available()"

发送数据

创建在进程间发送的 IPC 句柄，执行以下操作：

import torch
from dora.cuda import torch_to_ipc_buffer

torch_tensor = torch.tensor([1, 2, 3], dtype=torch.int64, device="cuda")
ipc_buffer, metadata = torch_to_ipc_buffer(torch_tensor)
node.send_output("latency", ipc_buffer, metadata)

接收数据

import pyarrow as pa
from dora import Node
from dora.cuda import ipc_buffer_to_ipc_handle, cudabuffer_to_torch

ctx = pa.cuda.context()
node = Node()
event = node.next() # 获取一个带有一个 torch 句柄的事件

ipc_handle = ipc_buffer_to_ipc_handle(event["value"])
cudabuffer = ctx.open_ipc_buffer(ipc_handle)
torch_tensor = cudabuffer_to_torch(cudabuffer, event["metadata"])  # 在 cuda