anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Feature] vLLM-Omni RDMA connector (#1019)

Browse Source 
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
This commit is contained in:
NATURE authored and GitHub committed 2026-02-24 11:13:52 +08:00
1 parent 7b02b856a6
commit 1589931fa6
20 files changed
+3645 -63
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

benchmarks/distributed/omni_connectors/README.md
+397
View File
@@ -0,0 +1,397 @@
# RDMA Test Configuration Guide
This document explains how to configure the RDMA environment and run tests for `MooncakeTransferEngineConnector`.
## Table of Contents
- [Docker Container Permissions](#docker-container-permissions)
- [Single-Node Testing](#single-node-testing)
- [Multi-Node Testing](#multi-node-testing)
- [Running Tests](#running-tests)
- [Cross-Node Testing](#cross-node-testing)
- [Troubleshooting](#troubleshooting)
---
## Docker Container Permissions
RDMA tests require access to InfiniBand/RoCE devices and system topology. Add the following permissions when running `docker run`.
### Option 1: Minimal Permissions (Recommended)
```bash
docker run -it \
--cap-add=SYS_PTRACE \
--cap-add=IPC_LOCK \
--security-opt seccomp=unconfined \
--network=host \
--device=/dev/infiniband \
-v /sys/class/infiniband:/sys/class/infiniband:ro \
your-image:tag
```
Parameter explanation:
- `--cap-add=SYS_PTRACE`: Allow reading system topology information
- `--cap-add=IPC_LOCK`: Allow memory locking (required for RDMA memory registration)
- `--security-opt seccomp=unconfined`: Disable seccomp restrictions
- `--network=host`: Use host network (required for RDMA)
- `--device=/dev/infiniband`: Mount InfiniBand devices
- `-v /sys/class/infiniband`: Mount IB device info (read-only)
### Option 2: Full Permissions (Quick but not recommended for production)
```bash
docker run -it \
--privileged \
--network=host \
your-image:tag
```
`--privileged` grants full host permissions. Suitable for quick testing but not recommended for production.
---
## Single-Node Testing
When running single-node tests (producer and consumer on the same machine), ensure they use the **same RDMA device**.
### Problem Background
InfiniBand devices use LID (Local Identifier) for routing. Different devices have different LIDs and cannot communicate directly. If no device is specified, Mooncake may assign different devices to connectors, causing handshake failures.
Common error:
```
[Handshake] Failed to modify QP to RTR, check mtu, gid, peer lid, peer qp num: Invalid argument [22]
```
### Solution
**Method 1: Set Environment Variable (Recommended)**
```bash
# List available RDMA devices
ibstat
# Select a device (e.g., mlx5_0)
export RDMA_DEVICE_NAME='mlx5_0'
# Run tests
pytest test_mooncake_transfer_engine_rdma.py -v -s
```
**Method 2: Use RoCE Devices**
If the system has RoCE devices (using IPv4 routing), the test code will automatically detect and prefer them. RoCE device GIDs start with `00:00:00:00:00:00:00:00:00:00:ff:ff` (IPv4-mapped).
**Method 3: Ensure MTU Consistency**
Make sure both endpoints use the same MTU:
```bash
# Check device MTU
ibstatus mlx5_0
```
---
## Multi-Node Testing
For multi-node tests, producer and consumer run on different machines connected via InfiniBand switch.
### Prerequisites
1. Both machines have Mooncake and RDMA drivers installed
2. Both machines are in the same InfiniBand subnet
3. Switch is properly configured
### Configuration
**Machine A (Producer):**
```bash
# Set RDMA host IP (InfiniBand interface IP)
export RDMA_TEST_HOST='10.0.0.1'
# Optional: Specify device
export RDMA_DEVICE_NAME='mlx5_0'
```
**Machine B (Consumer):**
```bash
# Set RDMA host IP
export RDMA_TEST_HOST='10.0.0.2'
# Optional: Specify device
export RDMA_DEVICE_NAME='mlx5_0'
```
### Verify Connectivity
```bash
# Ping IB interface
ping 10.0.0.2
# Test RDMA connectivity with ibping
# On Machine B (server)
ibping -S
# On Machine A (client)
ibping -G <Machine_B_GID>
```
---
## Running Tests
### Run All RDMA Tests (Single-Node, fast suite)
Slow tests (large payloads, stress, concurrency integrity) are marked `@pytest.mark.slow`. Use `-m "not slow"` to skip them in quick CI or local fast iteration.
```bash
cd tests/distributed/omni_connectors
# Fast suite only (excludes slow/stress tests)
pytest test_mooncake_transfer_engine_rdma.py test_mooncake_transfer_engine_buffer.py -v -s -m "not slow"
```
### Run Including Slow Tests
```bash
# Run ALL tests including slow/stress tests
pytest test_mooncake_transfer_engine_rdma.py test_mooncake_transfer_engine_buffer.py -v -s
# Run ONLY the slow/stress tests
pytest test_mooncake_transfer_engine_rdma.py test_mooncake_transfer_engine_buffer.py -v -s -m slow
```
### Run Buffer Management Tests
```bash
# Fast only
pytest test_mooncake_transfer_engine_buffer.py -v -s -m "not slow"
# Including allocator invariant tests (double-free, overlap, merge)
pytest test_mooncake_transfer_engine_buffer.py -v -s
```
### Run Specific Test Classes
```bash
# Basic connector tests
pytest test_mooncake_transfer_engine_rdma.py::TestBasicConnector -v -s
# End-to-end RDMA transfer tests
pytest test_mooncake_transfer_engine_rdma.py::TestEndToEnd -v -s
# Lifecycle & resource management tests
pytest test_mooncake_transfer_engine_rdma.py::TestLifecycle -v -s
# GPU memory pool tests (requires CUDA)
pytest test_mooncake_transfer_engine_rdma.py::TestGPUPool -v -s
# Stress / correctness tests (slow)
pytest test_mooncake_transfer_engine_rdma.py::TestStressCorrectness -v -s
```
### RDMA Environment Diagnostics
For quick diagnostics (device status, Mooncake availability, env vars, etc.),
see the [Troubleshooting section](../../../docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md#troubleshooting)
in the connector documentation.
---
## Cross-Node Testing
The `cross_node_mooncake_transfer_engine.py` script enables testing RDMA transfers between two separate physical machines. This script is **not** auto-discovered by `pytest` (it does not start with `test_`) — it must be run manually on each node.
### Prerequisites
1. Both machines have Mooncake installed
2. Both machines are connected via InfiniBand/RoCE switch
3. Firewall allows ZMQ ports (default: 15500, 15501)
4. Same RDMA device name on both nodes (if multiple devices exist)
### Running Cross-Node Tests
**On Machine A (Producer) — start first:**
```bash
cd benchmarks/distributed/omni_connectors/
# Optional: specify device if multiple exist
export RDMA_DEVICE_NAME='mlx5_0'
python cross_node_mooncake_transfer_engine.py \
--role producer \
--local-host <PRODUCER_IP> \
--remote-host <CONSUMER_IP> \
--tensor-size-mb 100 \
--num-transfers 3
```
**On Machine B (Consumer) — start after producer:**
```bash
cd benchmarks/distributed/omni_connectors/
export RDMA_DEVICE_NAME='mlx5_0'
python cross_node_mooncake_transfer_engine.py \
--role consumer \
--local-host <CONSUMER_IP> \
--remote-host <PRODUCER_IP> \
--tensor-size-mb 100 \
--num-transfers 3
```
### Transfer Modes
| Mode | Description | Example |
|------|-------------|---------|
| `copy` | Normal path — tensor copied to RDMA pool (default) | `--mode copy` |
| `zerocopy` | Zero-copy path — data created directly in RDMA pool | `--mode zerocopy` |
| `gpu` | GPU transfer — RDMA pool on GPU, uses GPUDirect | `--mode gpu --gpu-id 0` |
### Benchmark Mode
Skip MD5 verification and measure pure RDMA throughput:
```bash
# Producer
python cross_node_mooncake_transfer_engine.py \
--role producer \
--local-host <PRODUCER_IP> \
--remote-host <CONSUMER_IP> \
--tensor-size-mb 1024 \
--num-transfers 20 \
--benchmark
# Consumer
python cross_node_mooncake_transfer_engine.py \
--role consumer \
--local-host <CONSUMER_IP> \
--remote-host <PRODUCER_IP> \
--tensor-size-mb 1024 \
--num-transfers 20 \
--benchmark
```
### Cross-Node Test Options
| Option | Description | Default |
|--------|-------------|---------|
| `--role` | `producer` or `consumer` | Required |
| `--local-host` | Local RDMA IP address | Required |
| `--remote-host` | Remote RDMA IP address | Required |
| `--local-port` | Local ZMQ port for RDMA data | 15500 |
| `--remote-port` | Remote ZMQ port for RDMA data | 15500 |
| `--ctrl-port` | Control channel port | 15501 |
| `--tensor-size-mb` | Tensor size in MB | 100 |
| `--num-transfers` | Number of transfers | 3 |
| `--mode` | `copy`, `zerocopy`, or `gpu` | `copy` |
| `--gpu-id` | GPU ID for GPU mode | 0 |
| `--benchmark` | Skip MD5, pure performance test | off |
---
## Troubleshooting
### 1. "Failed to modify QP to RTR" Error
**Cause**: QP handshake failed, usually due to device configuration mismatch.
**Solution**:
```bash
# Force using the same device
export RDMA_DEVICE_NAME='mlx5_0'
```
### 2. "Mooncake TransferEngine is not available"
**Cause**: Mooncake not installed or import failed.
**Solution**:
```bash
# Check Mooncake installation
python -c "from mooncake.engine import TransferEngine; print('OK')"
# Reinstall if needed
pip install mooncake-transfer-engine
# Or using uv
uv pip install mooncake-transfer-engine
```
### 3. "Permission denied" accessing /dev/infiniband
**Cause**: Container lacks IB device access permissions.
**Solution**:
```bash
docker run --device=/dev/infiniband --cap-add=IPC_LOCK ...
```
### 4. Test Timeout
**Cause**: RDMA connection establishment failed or network latency.
**Solution**:
```bash
# Check network status
ibstat
ibstatus
```
### 5. GPU Test Failed "CUDA is not available"
**Cause**: CUDA environment not configured or GPU unavailable.
**Solution**:
```bash
# Check CUDA
python -c "import torch; print(torch.cuda.is_available())"
# Docker needs NVIDIA runtime
docker run --gpus all ...
```
---
## Environment Variables Reference
| Variable | Description | Example |
|----------|-------------|---------|
| `RDMA_DEVICE_NAME` | Specify RDMA device name | `mlx5_0` |
| `RDMA_TEST_HOST` | Specify test host IP | `10.0.0.1` |
| `MC_TE_METRIC` | Enable Mooncake metrics | `1` |
| `MC_IB_PCI_RELAXED_ORDERING` | Enable PCIe relaxed ordering | `1` |
---
## Test Files Overview
| File | Description | Auto-discovered by pytest |
|------|-------------|--------------------------|
| `test_mooncake_transfer_engine_rdma.py` | Integration tests for MooncakeTransferEngineConnector (basic, E2E, lifecycle, GPU) | Yes |
| `test_mooncake_transfer_engine_buffer.py` | Memory pool and buffer management unit tests | Yes |
| `cross_node_mooncake_transfer_engine.py` | Cross-node (multi-machine) testing script — run manually | No (filename does not start with `test_`) |
### test_mooncake_transfer_engine_rdma.py — Test Classes
| Test Class | Memory Pool | Marker | Description |
|------------|-------------|--------|-------------|
| `TestBasicConnector` | CPU | — | Initialization, put tensor/bytes/object, cleanup, pool exhaustion |
| `TestEndToEnd` | CPU | — | E2E RDMA transfer: tensor, bytes, object, zero-copy, large payload (100MB), mixed types, concurrency |
| `TestLifecycle` | CPU | — | Close, context manager, double-close safety |
| `TestGPUPool` | GPU | — | GPU pool init, put CPU/GPU tensor, GPU E2E transfer |
| `TestStressCorrectness` | CPU | `slow` | Concurrent put+get with MD5 integrity, bidirectional concurrency, edge cases (1-element tensor, empty bytes), 500MB payload, rapid alloc/free cycles |
### test_mooncake_transfer_engine_buffer.py — Test Classes
| Test Class | Marker | Description |
|------------|--------|-------------|
| `TestBufferAllocator` | — | Basic alloc/free, alignment, exhaustion/recovery, thread safety |
| `TestAllocatorInvariants` | `slow` | Double-free safety, overlap corruption detection, adjacent-block merging, fragmentation/defrag |
| `TestManagedBuffer` | — | Tensor views, context manager |
benchmarks/distributed/omni_connectors/cross_node_mooncake_transfer_engine.py
+644
View File
@@ -0,0 +1,644 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Cross-Node RDMA Test Script (Automated Version)
This script enables testing RDMA transfers between two separate machines.
Supports three transfer modes:
- copy: Normal path - tensor copied to RDMA pool (default)
- zerocopy: Zero-copy path - data created directly in RDMA pool
- gpu: GPU transfer - RDMA pool on GPU, uses GPUDirect
Usage:
# On Machine A (Producer) - start first:
python cross_node_mooncake_transfer_engine.py --role producer --local-host hostname_A --remote-host hostname_B
# On Machine B (Consumer) - start after producer:
python cross_node_mooncake_transfer_engine.py --role consumer --local-host hostname_B --remote-host hostname_A
# Zero-copy mode:
python cross_node_mooncake_transfer_engine.py --role producer ... --mode zerocopy
# GPU mode (requires GPUDirect RDMA support):
python cross_node_mooncake_transfer_engine.py --role producer ... --mode gpu --gpu-id 0
# Benchmark mode (skip random data generation and MD5 verification,
# measures pure RDMA throughput):
python cross_node_mooncake_transfer_engine.py --role producer ... --benchmark
Environment Variables:
RDMA_DEVICE_NAME: Specify RDMA device (e.g., mlx5_0)
MC_IB_PCI_RELAXED_ORDERING: Set to 1 to enable PCIe relaxed ordering
for higher RDMA throughput
"""
import argparse
import hashlib
import os
import sys
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Any
import msgspec
import torch
import zmq
# Add parent path for imports
sys.path.insert(
0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))
)
from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import (
ManagedBuffer,
MooncakeTransferEngineConnector,
TransferEngine,
)
def compute_md5(tensor: torch.Tensor) -> str:
"""Compute MD5 checksum of a tensor."""
if tensor.is_cuda:
tensor = tensor.cpu()
data = tensor.contiguous().view(torch.uint8).numpy().tobytes()
return hashlib.md5(data).hexdigest()
# Control channel message types
class CtrlMsg(msgspec.Struct):
"""Control channel message."""
msg_type: str # "READY", "TRANSFER", "ACK", "DONE", "ERROR"
request_id: str = ""
md5: str = ""
data_size: int = 0
error: str = ""
@dataclass
class TransferConfig:
"""Configuration for cross-node transfer test."""
local_host: str
remote_host: str
local_port: int
remote_port: int
ctrl_port: int
num_transfers: int
tensor_size_mb: int
mode: str # "copy", "zerocopy", "gpu"
gpu_id: int = 0
benchmark: bool = False # Skip MD5 verification for pure performance test
@dataclass
class TransferStats:
"""Statistics for transfer operations."""
success_count: int = 0
fail_count: int = 0
total_bytes: int = 0
elapsed_time: float = 0.0
@property
def throughput_mbps(self) -> float:
if self.elapsed_time > 0:
return (self.total_bytes / (1024 * 1024)) / self.elapsed_time
return 0.0
def print_summary(self, role: str):
print(f"\n{'=' * 60}")
print(f" {role.upper()} SUMMARY")
print(f" Successful: {self.success_count}/{self.success_count + self.fail_count}")
print(f" Failed: {self.fail_count}/{self.success_count + self.fail_count}")
print(f" Total: {self.total_bytes / (1024 * 1024):.2f} MB")
print(f" Time: {self.elapsed_time:.2f} s")
print(f" Throughput: {self.throughput_mbps:.2f} MB/s")
print(f"{'=' * 60}")
class CrossNodeTester(ABC):
"""Abstract base class for cross-node RDMA testing."""
def __init__(self, config: TransferConfig):
self.config = config
self.connector: MooncakeTransferEngineConnector | None = None
self.zmq_ctx: zmq.Context | None = None
self.ctrl_socket: zmq.Socket | None = None
self.stats = TransferStats()
def get_connector_config(self) -> dict:
"""Get connector configuration based on mode."""
pool_size = int(self.config.tensor_size_mb * 1.5) * 1024 * 1024
pool_size = max(pool_size, 128 * 1024 * 1024)
conn_config = {
"host": self.config.local_host,
"zmq_port": self.config.local_port,
"protocol": "rdma",
"memory_pool_size": pool_size,
}
# Set device based on mode
if self.config.mode == "gpu":
conn_config["memory_pool_device"] = f"cuda:{self.config.gpu_id}"
else:
conn_config["memory_pool_device"] = "cpu"
# RDMA device name from environment
device_name = os.environ.get("RDMA_DEVICE_NAME")
if device_name:
conn_config["device_name"] = device_name
print(f"[CONFIG] Using RDMA device: {device_name}")
return conn_config
def initialize(self):
"""Initialize connector and ZMQ context."""
print(f"[{self.role}] Initializing connector...")
conn_config = self.get_connector_config()
self.connector = MooncakeTransferEngineConnector(conn_config)
self.zmq_ctx = zmq.Context()
print(f"[{self.role}] Ready at {self.config.local_host}:{self.config.local_port}")
def cleanup(self):
"""Cleanup resources."""
if self.ctrl_socket:
self.ctrl_socket.close()
if self.zmq_ctx:
self.zmq_ctx.term()
if self.connector:
self.connector.close()
print(f"[{self.role}] Closed.")
@property
@abstractmethod
def role(self) -> str:
pass
@abstractmethod
def run(self):
pass
class Producer(CrossNodeTester):
"""Producer node - sends data to consumer."""
@property
def role(self) -> str:
return "PRODUCER"
def print_header(self):
print(f"\n{'=' * 60}")
print(f" PRODUCER MODE ({self.config.mode.upper()})")
print(f" Local: {self.config.local_host}:{self.config.local_port}")
print(f" Remote: {self.config.remote_host}:{self.config.remote_port}")
print(f" Control Port: {self.config.ctrl_port}")
print(f" Transfer Mode: {self.config.mode}")
if self.config.mode == "gpu":
print(f" GPU ID: {self.config.gpu_id}")
print(f"{'=' * 60}\n")
def setup_control_channel(self):
"""Setup ZMQ control channel as server (REP socket)."""
self.ctrl_socket = self.zmq_ctx.socket(zmq.REP)
self.ctrl_socket.bind(f"tcp://*:{self.config.ctrl_port}")
print(f"[PRODUCER] Control channel listening on port {self.config.ctrl_port}")
print("[PRODUCER] Waiting for consumer to connect...")
def wait_for_consumer(self) -> bool:
"""Wait for consumer READY signal."""
msg_data = self.ctrl_socket.recv()
msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
if msg.msg_type != "READY":
print(f"[PRODUCER] Unexpected message: {msg.msg_type}")
return False
print("[PRODUCER] Consumer connected!")
self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="ACK")))
return True
def create_test_data(self, transfer_idx: int) -> tuple[Any, str, int]:
"""
Create test data based on transfer mode.
Returns:
(data, md5, size) tuple
"""
num_elements = (self.config.tensor_size_mb * 1024 * 1024) // 4
data_size = num_elements * 4
if self.config.mode == "zerocopy":
# Zero-Copy Path: Allocate directly from connector's pool
offset = self.connector.allocator.alloc(data_size)
managed_buf = ManagedBuffer(self.connector.allocator, offset, data_size, self.connector.pool)
if self.config.benchmark:
# In benchmark mode, skip random data generation (use uninitialized memory)
return managed_buf, "", data_size
else:
# Fill buffer with random data using tensor view
tensor_view = managed_buf.as_tensor(dtype=torch.float32, shape=(num_elements,))
random_data = torch.randn(num_elements, dtype=torch.float32)
if tensor_view.is_cuda:
tensor_view.copy_(random_data.to(tensor_view.device))
else:
tensor_view.copy_(random_data)
md5 = compute_md5(tensor_view)
return managed_buf, md5, data_size
elif self.config.mode == "gpu":
# GPU Path: Create tensor on GPU
device = f"cuda:{self.config.gpu_id}"
if self.config.benchmark:
# In benchmark mode, use empty tensor (no random generation)
gpu_tensor = torch.empty(num_elements, dtype=torch.float32, device=device)
return gpu_tensor, "", data_size
else:
cpu_tensor = torch.randn(num_elements, dtype=torch.float32)
md5 = compute_md5(cpu_tensor)
gpu_tensor = cpu_tensor.to(device)
return gpu_tensor, md5, data_size
else:
# Copy Path (default): Create regular CPU tensor
if self.config.benchmark:
# In benchmark mode, use empty tensor (no random generation)
tensor = torch.empty(num_elements, dtype=torch.float32)
return tensor, "", data_size
else:
tensor = torch.randn(num_elements, dtype=torch.float32)
md5 = compute_md5(tensor)
return tensor, md5, data_size
def do_transfer(self, transfer_idx: int) -> bool:
"""Perform a single transfer."""
req_id = f"cross_node_transfer_{transfer_idx}"
if not self.config.benchmark:
print(f"\n[PRODUCER] Transfer {transfer_idx + 1}/{self.config.num_transfers}")
# Create test data
t0 = time.time()
data, md5, data_size = self.create_test_data(transfer_idx)
t_create = time.time() - t0
if not self.config.benchmark:
print(f" Mode: {self.config.mode}")
print(f" Size: {self.config.tensor_size_mb} MB")
if md5:
print(f" MD5: {md5[:16]}...")
print(f" Create time: {t_create * 1000:.1f} ms")
# Put data
t1 = time.time()
success, size, metadata = self.connector.put("producer", "consumer", req_id, data)
t_put = time.time() - t1
if not success:
print(" [FAIL] Put failed")
return False
if not self.config.benchmark:
print(f" [OK] Put successful, {size} bytes ({t_put * 1000:.1f} ms)")
# Wait for consumer to request transfer info
msg_data = self.ctrl_socket.recv()
msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
if msg.msg_type != "READY":
print(f" [ERROR] Unexpected message: {msg.msg_type}")
return False
# Send transfer metadata to consumer
transfer_msg = CtrlMsg(
msg_type="TRANSFER",
request_id=req_id,
md5=md5,
data_size=data_size,
)
self.ctrl_socket.send(msgspec.msgpack.encode(transfer_msg))
# Wait for consumer ACK (this includes RDMA transfer time)
t2 = time.time()
msg_data = self.ctrl_socket.recv()
t_rdma = time.time() - t2
msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
success = msg.msg_type == "ACK"
if success:
if not self.config.benchmark:
print(f" [OK] RDMA transfer complete ({t_rdma * 1000:.1f} ms)")
self.stats.success_count += 1
self.stats.total_bytes += size
else:
print(f" [WARN] Consumer reported error: {msg.error}")
self.stats.fail_count += 1
# Send ACK to allow consumer to continue
self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="ACK")))
# Cleanup buffer
self.connector.cleanup(req_id)
return success
def run(self):
"""Run the producer."""
self.print_header()
self.initialize()
self.setup_control_channel()
try:
if not self.wait_for_consumer():
return
if self.config.benchmark:
print(
f"[BENCHMARK] Running {self.config.num_transfers} "
f"transfers of {self.config.tensor_size_mb} MB each..."
)
start_time = time.time()
for i in range(self.config.num_transfers):
self.do_transfer(i)
if self.config.benchmark and (i + 1) % 10 == 0:
elapsed = time.time() - start_time
current_throughput = (self.stats.total_bytes / (1024 * 1024)) / elapsed
print(f" Progress: {i + 1}/{self.config.num_transfers}, Throughput: {current_throughput:.2f} MB/s")
self.stats.elapsed_time = time.time() - start_time
self.stats.print_summary("PRODUCER")
# Wait for final consumer message and send DONE
self.ctrl_socket.recv()
self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="DONE")))
finally:
self.cleanup()
class Consumer(CrossNodeTester):
"""Consumer node - receives data from producer."""
@property
def role(self) -> str:
return "CONSUMER"
def print_header(self):
print(f"\n{'=' * 60}")
print(f" CONSUMER MODE ({self.config.mode.upper()})")
print(f" Local: {self.config.local_host}:{self.config.local_port}")
print(f" Remote: {self.config.remote_host}:{self.config.remote_port}")
print(f" Control Port: {self.config.ctrl_port}")
print(f" Transfer Mode: {self.config.mode}")
if self.config.mode == "gpu":
print(f" GPU ID: {self.config.gpu_id}")
print(f"{'=' * 60}\n")
def setup_control_channel(self):
"""Setup ZMQ control channel as client (REQ socket)."""
self.ctrl_socket = self.zmq_ctx.socket(zmq.REQ)
ctrl_addr = f"tcp://{self.config.remote_host}:{self.config.ctrl_port}"
print(f"[CONSUMER] Connecting to producer control channel at {ctrl_addr}...")
self.ctrl_socket.connect(ctrl_addr)
def connect_to_producer(self) -> bool:
"""Connect to producer and send READY signal."""
self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="READY")))
msg_data = self.ctrl_socket.recv()
msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
if msg.msg_type != "ACK":
print(f"[CONSUMER] Unexpected response: {msg.msg_type}")
return False
print("[CONSUMER] Connected to producer! Starting transfers...")
return True
def do_transfer(self, transfer_idx: int) -> bool:
"""Perform a single transfer."""
if not self.config.benchmark:
print(f"\n[CONSUMER] Transfer {transfer_idx + 1}/{self.config.num_transfers}")
# Request next transfer info
self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="READY")))
msg_data = self.ctrl_socket.recv()
msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
if msg.msg_type == "DONE":
print("[CONSUMER] Producer signaled completion")
return False
if msg.msg_type != "TRANSFER":
print(f"[CONSUMER] Unexpected message: {msg.msg_type}")
return False
req_id = msg.request_id
expected_md5 = msg.md5
data_size = msg.data_size
num_elements = data_size // 4
if not self.config.benchmark:
print(f" Request ID: {req_id}")
if expected_md5:
print(f" Expected MD5: {expected_md5[:16]}...")
print(f" Data Size: {data_size / (1024 * 1024):.2f} MB")
# Build metadata for get
metadata = {
"request_id": req_id,
"source_host": self.config.remote_host,
"source_port": self.config.remote_port,
"data_size": data_size,
"dtype": "float32",
"shape": [num_elements],
"is_fast_path": True,
}
if not self.config.benchmark:
print(f" [INFO] Requesting from {self.config.remote_host}:{self.config.remote_port}")
# Get data with timing
t0 = time.time()
result = self.connector.get("producer", "consumer", req_id, metadata)
t_get = time.time() - t0
response_msg = CtrlMsg(msg_type="ERROR", error="Get failed")
if result is not None:
recv_buffer, recv_size = result
if not self.config.benchmark:
print(f" [OK] Get successful, {recv_size} bytes ({t_get * 1000:.1f} ms)")
if isinstance(recv_buffer, ManagedBuffer):
# In benchmark mode, skip MD5 verification
if self.config.benchmark or not expected_md5:
response_msg = CtrlMsg(msg_type="ACK")
self.stats.success_count += 1
self.stats.total_bytes += recv_size
else:
# Verify data
t1 = time.time()
reconstructed = recv_buffer.as_tensor(dtype=torch.float32, shape=(num_elements,))
recv_md5 = compute_md5(reconstructed)
t_md5 = time.time() - t1
print(f" MD5: {recv_md5[:16]}... ({t_md5 * 1000:.1f} ms)")
if recv_md5 == expected_md5:
print(" [PASS] MD5 checksum verified!")
response_msg = CtrlMsg(msg_type="ACK")
self.stats.success_count += 1
self.stats.total_bytes += recv_size
else:
print(" [FAIL] MD5 mismatch!")
response_msg = CtrlMsg(msg_type="ERROR", error="MD5 mismatch")
self.stats.fail_count += 1
recv_buffer.release()
else:
response_msg = CtrlMsg(msg_type="ACK")
self.stats.success_count += 1
self.stats.total_bytes += recv_size
else:
print(" [FAIL] Get failed")
self.stats.fail_count += 1
# Send response to producer
self.ctrl_socket.send(msgspec.msgpack.encode(response_msg))
# Wait for ACK
self.ctrl_socket.recv()
return response_msg.msg_type == "ACK"
def run(self):
"""Run the consumer."""
self.print_header()
self.initialize()
self.setup_control_channel()
try:
if not self.connect_to_producer():
return
if self.config.benchmark:
print(
f"[BENCHMARK] Running {self.config.num_transfers} "
f"transfers of {self.config.tensor_size_mb} MB each..."
)
start_time = time.time()
for i in range(self.config.num_transfers):
if not self.do_transfer(i):
break
if self.config.benchmark and (i + 1) % 10 == 0:
elapsed = time.time() - start_time
current_throughput = (self.stats.total_bytes / (1024 * 1024)) / elapsed
print(f" Progress: {i + 1}/{self.config.num_transfers}, Throughput: {current_throughput:.2f} MB/s")
self.stats.elapsed_time = time.time() - start_time
self.stats.print_summary("CONSUMER")
# Send final READY and wait for DONE
self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="READY")))
self.ctrl_socket.recv()
finally:
self.cleanup()
def main():
parser = argparse.ArgumentParser(
description="Cross-Node RDMA Test (Automated)",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Transfer Modes:
copy - Normal path: tensor copied to RDMA pool (default)
zerocopy - Zero-copy path: data created directly in RDMA pool
gpu - GPU transfer: RDMA pool on GPU, uses GPUDirect
Examples:
# Copy mode (default):
python cross_node_mooncake_transfer_engine.py --role producer \
--local-host hostA --remote-host hostB
# Zero-copy mode:
python cross_node_mooncake_transfer_engine.py --role producer \
--local-host hostA --remote-host hostB --mode zerocopy
# GPU mode:
python cross_node_mooncake_transfer_engine.py --role producer \
--local-host hostA --remote-host hostB --mode gpu --gpu-id 0
# Benchmark mode (skip MD5, measure pure RDMA performance):
python cross_node_mooncake_transfer_engine.py --role producer \
--local-host hostA --remote-host hostB --benchmark
# With specific RDMA device:
RDMA_DEVICE_NAME=mlx5_0 python cross_node_mooncake_transfer_engine.py --role producer ...
""",
)
parser.add_argument(
"--role", required=True, choices=["producer", "consumer"], help="Role: producer (sends) or consumer (receives)"
)
parser.add_argument("--local-host", required=True, help="Local hostname or IP address")
parser.add_argument("--remote-host", required=True, help="Remote hostname or IP address")
parser.add_argument("--local-port", type=int, default=15500, help="Local ZMQ port for RDMA data (default: 15500)")
parser.add_argument("--remote-port", type=int, default=15500, help="Remote ZMQ port for RDMA data (default: 15500)")
parser.add_argument("--ctrl-port", type=int, default=15501, help="Control channel port (default: 15501)")
parser.add_argument("--num-transfers", type=int, default=20, help="Number of transfers to perform (default: 3)")
parser.add_argument("--tensor-size-mb", type=int, default=100, help="Tensor size in MB (default: 100)")
parser.add_argument(
"--mode",
choices=["copy", "zerocopy", "gpu"],
default="copy",
help="Transfer mode: copy, zerocopy, or gpu (default: copy)",
)
parser.add_argument("--gpu-id", type=int, default=0, help="GPU ID for GPU mode (default: 0)")
parser.add_argument(
"--benchmark", action="store_true", help="Benchmark mode: skip MD5 verification for pure performance test"
)
args = parser.parse_args()
# Check Mooncake
if TransferEngine is None:
print("[ERROR] Mooncake TransferEngine is not available.")
print("Install with: pip install mooncake")
sys.exit(1)
# Check CUDA for GPU mode
if args.mode == "gpu":
if not torch.cuda.is_available():
print("[ERROR] CUDA is not available but GPU mode was requested.")
sys.exit(1)
if args.gpu_id >= torch.cuda.device_count():
print(f"[ERROR] GPU {args.gpu_id} not available. Found {torch.cuda.device_count()} GPUs.")
sys.exit(1)
print(f"[INFO] Using GPU {args.gpu_id}: {torch.cuda.get_device_name(args.gpu_id)}")
config = TransferConfig(
local_host=args.local_host,
remote_host=args.remote_host,
local_port=args.local_port,
remote_port=args.remote_port,
ctrl_port=args.ctrl_port,
num_transfers=args.num_transfers,
tensor_size_mb=args.tensor_size_mb,
mode=args.mode,
gpu_id=args.gpu_id,
benchmark=args.benchmark,
)
if args.role == "producer":
tester = Producer(config)
else:
tester = Consumer(config)
tester.run()
if __name__ == "__main__":
main()