diff --git a/docs/source/design/transfer-engine/efa_transport.md b/docs/source/design/transfer-engine/efa_transport.md new file mode 100644 index 00000000..b6fee161 --- /dev/null +++ b/docs/source/design/transfer-engine/efa_transport.md @@ -0,0 +1,353 @@ +# AWS EFA Transport for Mooncake + +This document describes how to build and use Mooncake with AWS Elastic Fabric Adapter (EFA) support using libfabric. + +## Prerequisites + +### 1. AWS EFA Driver and libfabric + +EFA driver and libfabric should be pre-installed on AWS instances with EFA support (e.g., p6-b200.48xlarge, p5e.48xlarge, p4d.24xlarge). + +Verify installation: +```bash +# Check EFA devices +fi_info -p efa + +# Verify libfabric location +ls /opt/amazon/efa/lib/libfabric.so +ls /opt/amazon/efa/include/rdma/fabric.h +``` + +If not installed, follow [AWS EFA documentation](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/efa-start.html). + +### 2. Build Dependencies + +```bash +# Ubuntu/Debian +sudo apt-get update +sudo apt-get install -y \ + build-essential \ + cmake \ + git \ + libgflags-dev \ + libgoogle-glog-dev \ + libjsoncpp-dev \ + libnuma-dev \ + libibverbs-dev \ + libboost-all-dev \ + libcurl4-openssl-dev \ + libyaml-cpp-dev \ + libgtest-dev \ + pybind11-dev \ + python3-dev + +# Install yalantinglibs (required) +cd /tmp +git clone https://github.com/alibaba/yalantinglibs.git +cd yalantinglibs +mkdir build && cd build +cmake .. -DCMAKE_INSTALL_PREFIX=/usr/local +make -j$(nproc) +sudo make install +``` + +## Building Mooncake with EFA Support + +### 1. Clone the Repository + +```bash +git clone https://github.com/kvcache-ai/Mooncake.git +cd Mooncake +git submodule update --init --recursive +``` + +### 2. Build with EFA Enabled + +```bash +mkdir build && cd build + +cmake .. \ + -DUSE_EFA=ON \ + -DWITH_TE=ON \ + -DWITH_STORE=ON \ + -DBUILD_UNIT_TESTS=ON \ + -DCMAKE_BUILD_TYPE=RelWithDebInfo + +make -j$(nproc) +``` + +### 3. Install Python Package + +```bash +# Copy built modules to wheel directory +cp mooncake-integration/engine.cpython-*.so ../mooncake-wheel/mooncake/ +cp mooncake-asio/libasio.so ../mooncake-wheel/mooncake/ + +# Install with pip +pip install -e ../mooncake-wheel --no-build-isolation +``` + +## Verification + +Test EFA transport initialization: + +```python +from mooncake.engine import TransferEngine + +te = TransferEngine() +result = te.initialize('127.0.0.1', 'P2PHANDSHAKE', 'efa', '') +print(f'Initialize result: {result}') # Should be 0 + +# You should see logs like: +# EFA device (libfabric): rdmap79s0, domain: rdmap79s0-rdm, provider: efa +``` + +## Usage with vLLM + +### Prefill Instance + +```bash +VLLM_MOONCAKE_BOOTSTRAP_PORT=8998 \ +vllm serve -tp 8 \ + --port 8010 \ + --trust-remote-code \ + --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_producer","kv_connector_extra_config":{"mooncake_protocol":"efa"}}' +``` + +### Decode Instance + +```bash +vllm serve -tp 8 \ + --port 8020 \ + --trust-remote-code \ + --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_consumer","kv_connector_extra_config":{"mooncake_protocol":"efa"}}' +``` + +## Unit Tests + +Run the EFA transport unit tests (requires EFA hardware): + +```bash +./build/mooncake-transfer-engine/tests/efa_transport_test +``` + +The test suite includes: + +| Test | Description | +|------|-------------| +| `InstallTransport` | Verify EFA transport installation | +| `LoopbackWrite` | Loopback write operation | +| `WriteAndRead` | Write then read with data integrity check | +| `MultiWrite` | Batch write (16 requests) | +| `StressMultipleBatches` | Stress test (20 batches x 8 requests) | + +You can also run all unit tests via CTest: + +```bash +cd build && ctest --output-on-failure +``` + +Environment variables for test configuration: + +```bash +export MC_METADATA_SERVER=P2PHANDSHAKE # default +export MC_LOCAL_SERVER_NAME=127.0.0.1:12345 # default +``` + +## Performance Benchmark + +Use `transfer_engine_bench` to measure EFA transport throughput between two nodes. + +### Target Node (receiver) + +```bash +./build/mooncake-transfer-engine/example/transfer_engine_bench \ + --mode=target \ + --protocol=efa \ + --metadata_server=P2PHANDSHAKE +``` + +### Initiator Node (sender) + +```bash +./build/mooncake-transfer-engine/example/transfer_engine_bench \ + --mode=initiator \ + --protocol=efa \ + --metadata_server=P2PHANDSHAKE \ + --segment_id=: \ + --operation=write \ + --duration=10 \ + --threads=8 \ + --block_size=65536 \ + --batch_size=128 \ + --buffer_size=1073741824 \ + --report_unit=GB +``` + +Replace `:` with the target node's address shown in the target's startup log (e.g., `ip-172-31-29-226:12345`). + +### Key Parameters + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `--block_size` | 65536 | Bytes per transfer request | +| `--batch_size` | 128 | Requests per batch | +| `--threads` | 12 | Concurrent submission threads | +| `--buffer_size` | 1 GB | Total buffer size | +| `--duration` | 10 | Test duration in seconds | +| `--operation` | read | `read` or `write` | +| `--report_unit` | GB | `GB\|GiB\|Gb\|MB\|MiB\|Mb` | + +### Benchmark Results + +Tested on two p6-b200.48xlarge instances (8 EFA devices each, 8×400 Gbps) in the same AWS placement group. + +#### Optimized Results + +With tuned parameters (`MC_SLICE_SIZE=262144`): + +| Operation | Throughput | Configuration | +|-----------|-----------|---------------| +| **Write** | **167.63 GB/s** | threads=48, block_size=128KB, batch_size=128, MC_SLICE_SIZE=256KB | +| **Read** | **171.89 GB/s** | threads=48, block_size=128KB, batch_size=128, MC_SLICE_SIZE=256KB | + +#### Parameter Tuning Results + +The following table shows how different parameters affect write throughput: + +| block_size | threads | batch_size | MC_SLICE_SIZE | Throughput | +|-----------|---------|------------|---------------|-----------| +| 64KB | 8 | 128 | default (64KB) | 69.47 GB/s | +| 256KB | 8 | 128 | default | 70.09 GB/s | +| 64KB | 16 | 128 | default | 78.80 GB/s | +| 64KB | 32 | 256 | default | 87.65 GB/s | +| 64KB | 64 | 256 | default | 85.72 GB/s | +| 128KB | 32 | 128 | default | 92.33 GB/s | +| 128KB | 32 | 128 | 128KB | 152.26 GB/s | +| 128KB | 32 | 128 | 256KB | 156.18 GB/s | +| 128KB | 48 | 128 | 256KB | **160.34 GB/s** | +| 128KB | 64 | 128 | 256KB | 158.82 GB/s | + +Key findings: +- **MC_SLICE_SIZE** is the most impactful tuning parameter — increasing from default 64KB to 256KB nearly **doubles** throughput (92→160 GB/s) +- **block_size=128KB** outperforms 64KB by ~10-15% +- **threads=48** is optimal for 8 EFA devices; 64 threads shows slight diminishing returns +- **batch_size=128** is sufficient; increasing to 256+ causes "Cannot select device" errors at higher thread counts + +#### Cross-Transport Comparison + +| Transport | Throughput | Per-NIC Bandwidth | Notes | +|-----------|-----------|-------------------|-------| +| **EFA (tuned)** | **168-172 GB/s** | ~207-214 Gbps × 8 NICs | MC_SLICE_SIZE=256KB, threads=48 | +| **EFA (default)** | **69.47 GB/s** | ~86 Gbps × 8 NICs | Default parameters | +| TCP (iperf3 baseline) | 9.5 GB/s | 76 Gbps total | Kernel TCP stack, 8 parallel streams | +| TCP (Mooncake) | 0.11 GB/s | — | Mooncake TCP transport, unoptimized for throughput | + +**EFA (tuned) vs TCP**: EFA delivers **17.7x** the raw TCP bandwidth by bypassing the kernel network stack. + +**EFA vs RoCE RDMA**: On comparable 8×400 Gbps RoCE networks, Mooncake's RDMA transport achieves ~190 GB/s. Tuned EFA reaches **~88%** of RoCE performance, demonstrating that proper parameter tuning can largely close the gap between SRD-based EFA and hardware-offloaded RDMA. + +### Tuning Tips + +- **Set `MC_SLICE_SIZE=262144` (256KB)** — this is the single most important tuning knob, nearly doubling throughput from defaults +- Increase `--threads` to 32-48 to saturate multiple EFA devices (6 threads per device is a good starting point) +- Use `--block_size=131072` (128KB) for optimal per-request efficiency +- Keep `--batch_size=128`; higher values may cause device selection failures with many threads +- Allocate buffers on both NUMA nodes for balanced NIC utilization (the bench tool does this by default) +- Avoid `--block_size=256KB` or larger with many threads — this can trigger "Cannot select device" errors due to buffer boundary alignment across 8 EFA devices + +## Technical Details + +### Why libfabric instead of ibverbs? + +AWS EFA exposes RDMA-like devices through the ibverbs interface, but does not support the full ibverbs API. Specifically: +- Queue Pair (QP) creation fails with "Operation not supported" (error 95) +- EFA requires using libfabric's `FI_EP_RDM` (Reliable Datagram Message) endpoint type + +### EFA Transport Architecture + +``` +┌─────────────────────────────────────────────────────┐ +│ EfaTransport │ +├─────────────────────────────────────────────────────┤ +│ EfaContext (per device) │ +│ ├── fi_info (fabric info) │ +│ ├── fid_fabric (fabric handle) │ +│ ├── fid_domain (protection domain) │ +│ ├── fid_av (address vector for peer lookup) │ +│ ├── fid_cq (completion queues) │ +│ └── fid_mr (memory regions) │ +├─────────────────────────────────────────────────────┤ +│ EfaEndpoint (per connection) │ +│ ├── fid_ep (RDM endpoint) │ +│ ├── fi_addr_t (peer address) │ +│ └── local_addr (local endpoint address) │ +└─────────────────────────────────────────────────────┘ +``` + +### Thread Safety + +The EFA transport requests `FI_THREAD_SAFE` from the libfabric provider and adds per-endpoint spinlocks to serialize `fi_write` calls. This is necessary because: + +- Multiple submission threads may route slices to the same endpoint concurrently +- libfabric RDM endpoints default to `FI_THREAD_UNSPEC` (no thread safety guarantees) +- Concurrent `fi_write` without serialization corrupts provider internals, causing completions to silently vanish + +CQ completion queues are polled by dedicated worker threads (one per EFA device) that run independently of submission threads. + +### EFA vs RoCE RDMA + +| Feature | EFA (libfabric SRD) | RoCE (ibverbs) | +|---------|--------------------|--------------------| +| Protocol | Scalable Reliable Datagram | RDMA over Converged Ethernet | +| Endpoint type | `FI_EP_RDM` (message-based) | Queue Pairs (true RDMA) | +| Write operation | Software-emulated via messages + ACKs | Hardware-offloaded one-sided RDMA | +| CPU overhead | Moderate (provider processes ACKs) | Minimal (NIC handles everything) | +| Throughput (8×400G) | ~170 GB/s (tuned) | ~190 GB/s | +| AWS availability | All EFA-enabled instances | Not available on AWS | + +### Supported AWS Instance Types + +- p6-b200.48xlarge (8 EFA devices, `rdmap*` naming) +- p5e.48xlarge (16 EFA devices, `rdmap*` naming) +- p4d.24xlarge (4 EFA devices) +- Other EFA-enabled instances + +Use `fi_info -p efa` to list available EFA devices on your instance. + +## Troubleshooting + +### No EFA devices found + +``` +EfaTransport: No EFA devices found +``` + +Solution: Verify EFA is available with `fi_info -p efa` + +### Permission denied + +``` +fi_fabric failed: Permission denied +``` + +Solution: Ensure proper permissions or run with sudo for testing + +### libfabric not found + +``` +cannot find -lfabric +``` + +Solution: Verify `/opt/amazon/efa/lib` is in the library path: +```bash +export LD_LIBRARY_PATH=/opt/amazon/efa/lib:$LD_LIBRARY_PATH +``` + +### Workers hang under high concurrency + +If `transfer_engine_bench` hangs with some workers never completing: + +1. **Ensure both nodes are running the same build** — the CQ backpressure and thread-safety fixes must be present on both sides +2. **Reduce concurrency** to verify basic connectivity: `--threads=1 --batch_size=16` +3. **Check CQ poller threads**: logs should show "Started N CQ polling worker threads" where N matches the number of EFA devices diff --git a/docs/source/design/transfer-engine/index.md b/docs/source/design/transfer-engine/index.md index 98146661..119b3775 100644 --- a/docs/source/design/transfer-engine/index.md +++ b/docs/source/design/transfer-engine/index.md @@ -11,7 +11,7 @@ Mooncake Transfer Engine is a high-performance, zero-copy data transfer library As shown in the diagram, each specific client corresponds to a `TransferEngine`, which not only includes a RAM Segment but also integrates management for high-speed transfers across multiple threads and network cards. The RAM Segment, in principle, corresponds to the entire virtual address space of this `TransferEngine`, but in reality, only parts of it (known as a `Buffer`) are registered for (GPUDirect) RDMA Read/Write. Each Buffer can have separate permissions (corresponding to RDMA `rkey`, etc.) and network card affinity (e.g., preferred NICs for different types of memory). -Mooncake Transfer Engine provides interfaces through the `TransferEngine` class (located in `mooncake-transfer-engine/include/transfer_engine.h`), where the specific data transfer functions for different backends are implemented by the `Transport` class, currently supporting `TcpTransport`, `RdmaTransport`, `NVMeoFTransport`, `NvlinkTransport`, `IntraNodeNvlinkTransport`, and `HipTransport`. +Mooncake Transfer Engine provides interfaces through the `TransferEngine` class (located in `mooncake-transfer-engine/include/transfer_engine.h`), where the specific data transfer functions for different backends are implemented by the `Transport` class, currently supporting `TcpTransport`, `RdmaTransport`, `EfaTransport`, `NVMeoFTransport`, `NvlinkTransport`, `IntraNodeNvlinkTransport`, and `HipTransport`. ### Segment Segment represents a collection of source address ranges and target address ranges available during the data transfer process in Transfer Engine. That is, all local and remote addresses involved in `BatchTransfer` requests must be within the valid segment range. Transfer Engine supports the following two types of Segments. @@ -155,7 +155,7 @@ The following video shows a normal run as described above, with the Target on th ![transfer-engine-running](../../image/transfer-engine-running.gif) ## Transfer Engine C/C++ API -Transfer Engine provides interfaces through the `TransferEngine` class (located in `mooncake-transfer-engine/include/transfer_engine.h`), where the specific data transfer functions for different backends are implemented by the `Transport` class, currently supporting `TcpTransport`, `RdmaTransport`, `NVMeoFTransport`, `NvlinkTransport` (for NVIDIA GPUs), `IntraNodeNvlinkTransport` (for NVIDIA GPUs), and `HipTransport` (for AMD GPUs). +Transfer Engine provides interfaces through the `TransferEngine` class (located in `mooncake-transfer-engine/include/transfer_engine.h`), where the specific data transfer functions for different backends are implemented by the `Transport` class, currently supporting `TcpTransport`, `RdmaTransport`, `EfaTransport` (for AWS EFA), `NVMeoFTransport`, `NvlinkTransport` (for NVIDIA GPUs), `IntraNodeNvlinkTransport` (for NVIDIA GPUs), and `HipTransport` (for AMD GPUs). For a complete C++ API reference, see [Transfer Engine C++ API Reference](cpp-api.md). @@ -311,6 +311,14 @@ For advanced users, TransferEngine provides the following advanced runtime optio cpp-api :::: +## EFA Transport (AWS) + +:::{toctree} +:maxdepth: 1 + +efa_transport +::: + ## Ascend Transport Component :::{toctree} diff --git a/mooncake-common/common.cmake b/mooncake-common/common.cmake index 4855061c..318b2464 100644 --- a/mooncake-common/common.cmake +++ b/mooncake-common/common.cmake @@ -68,6 +68,29 @@ option(USE_UBSHMEM "option for using ascend npu with shmem" OFF) option(USE_ASCEND_HETEROGENEOUS "option for transferring between ascend npu and gpu" OFF) option(USE_MNNVL "option for using Multi-Node NVLink transport" OFF) option(USE_CXL "option for using CXL protocol" OFF) +option(USE_EFA "option for using AWS EFA transport" OFF) + +if (USE_EFA) + # Find libfabric headers and library; default to AWS EFA installer path + find_path(LIBFABRIC_INCLUDE_DIR rdma/fabric.h + HINTS /opt/amazon/efa/include + PATH_SUFFIXES include) + find_library(LIBFABRIC_LIBRARY fabric + HINTS /opt/amazon/efa/lib + PATH_SUFFIXES lib lib64) + + if (NOT LIBFABRIC_INCLUDE_DIR OR NOT LIBFABRIC_LIBRARY) + message(FATAL_ERROR "libfabric not found. Install AWS EFA or set LIBFABRIC_INCLUDE_DIR/LIBFABRIC_LIBRARY.") + endif() + + get_filename_component(LIBFABRIC_LIB_DIR ${LIBFABRIC_LIBRARY} DIRECTORY) + include_directories(${LIBFABRIC_INCLUDE_DIR}) + link_directories(${LIBFABRIC_LIB_DIR}) + add_compile_definitions(USE_EFA) + message(STATUS "AWS EFA (libfabric) transport is enabled") + message(STATUS " libfabric include: ${LIBFABRIC_INCLUDE_DIR}") + message(STATUS " libfabric library: ${LIBFABRIC_LIBRARY}") +endif() option(USE_ETCD "option for enable etcd as metadata server" OFF) option(USE_ETCD_LEGACY "option for enable etcd based on etcd-cpp-api-v3" OFF) option(USE_REDIS "option for enable redis as metadata server" OFF) diff --git a/mooncake-integration/CMakeLists.txt b/mooncake-integration/CMakeLists.txt index 10c9284e..c677e2ff 100644 --- a/mooncake-integration/CMakeLists.txt +++ b/mooncake-integration/CMakeLists.txt @@ -39,6 +39,11 @@ if (WITH_TE) INSTALL_RPATH "$ORIGIN" ) + # Propagate EFA compile definition to engine target + if(USE_EFA) + target_compile_definitions(engine PRIVATE USE_EFA) + endif() + target_link_libraries(engine PRIVATE $ ) diff --git a/mooncake-integration/transfer_engine/transfer_engine_py.cpp b/mooncake-integration/transfer_engine/transfer_engine_py.cpp index b193302e..da45791d 100644 --- a/mooncake-integration/transfer_engine/transfer_engine_py.cpp +++ b/mooncake-integration/transfer_engine/transfer_engine_py.cpp @@ -158,12 +158,30 @@ int TransferEnginePy::initializeExt(const char *local_hostname, const char *protocol, const char *device_name, const char *metadata_type) { - (void)(protocol); + std::string proto = protocol ? std::string(protocol) : ""; std::string conn_string = buildConnString(metadata_type, metadata_server); auto device_name_safe = device_name ? std::string(device_name) : ""; auto device_filter = buildDeviceFilter(device_name_safe); + +#ifdef USE_EFA + // When using EFA protocol, we still need topology discovery but won't + // auto-install RDMA + bool use_efa = (proto == "efa"); + // Disable auto_discover to prevent RDMA transport installation, we'll + // install EFA manually + engine_ = std::make_unique(false, device_filter); + // Manually discover topology for EFA to populate device list + if (use_efa) { + engine_->getLocalTopology()->discover(device_filter); + LOG(INFO) << "Topology discovery complete for EFA. Found " + << engine_->getLocalTopology()->getHcaList().size() + << " devices."; + } +#else engine_ = std::make_unique(true, device_filter); +#endif + if (getenv("MC_LEGACY_RPC_PORT_BINDING")) { auto hostname_port = parseHostNameWithPort(local_hostname); int ret = @@ -176,6 +194,20 @@ int TransferEnginePy::initializeExt(const char *local_hostname, if (ret) return -1; } +#ifdef USE_EFA + // Install EFA transport when protocol is "efa" + if (use_efa) { + LOG(INFO) + << "Installing EFA transport as requested by protocol parameter"; + auto transport = engine_->installTransport("efa", nullptr); + if (!transport) { + LOG(ERROR) << "Failed to install EFA transport"; + return -1; + } + LOG(INFO) << "EFA transport installed successfully"; + } +#endif + free_list_.resize(kSlabSizeKBTabLen); return 0; } diff --git a/mooncake-transfer-engine/example/efa_latency_bench.png b/mooncake-transfer-engine/example/efa_latency_bench.png new file mode 100644 index 00000000..34693c07 Binary files /dev/null and b/mooncake-transfer-engine/example/efa_latency_bench.png differ diff --git a/mooncake-transfer-engine/example/efa_latency_bench.py b/mooncake-transfer-engine/example/efa_latency_bench.py new file mode 100755 index 00000000..42f288e2 --- /dev/null +++ b/mooncake-transfer-engine/example/efa_latency_bench.py @@ -0,0 +1,454 @@ +#!/usr/bin/env python3 +# Copyright 2024 KVCache.AI +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +EFA Latency vs Cache Size Benchmark + +Runs transfer_engine_bench on two machines via SSH to measure throughput +for multiple EFA configurations. For each cache size point, an independent +benchmark run measures throughput, then latency = cache_size / throughput +is computed. Running separate measurements per point captures natural +throughput variation between runs. + +Produces a chart similar to image/transfer-engine-performance.png. + +Usage: + python3 efa_latency_bench.py + python3 efa_latency_bench.py --target_host=HOST_A --initiator_host=HOST_B + python3 efa_latency_bench.py --output=my_chart.png --duration=10 +""" + +import argparse +import os +import re +import subprocess +import sys +import time + + +def parse_args(): + parser = argparse.ArgumentParser( + description="EFA Latency vs Cache Size Benchmark" + ) + parser.add_argument( + "--target_host", + default="ip-172-31-22-204", + help="Hostname or IP of the target machine (default: ip-172-31-22-204)", + ) + parser.add_argument( + "--initiator_host", + default="ip-172-31-26-160", + help="Hostname or IP of the initiator machine (default: ip-172-31-26-160)", + ) + parser.add_argument( + "--build_dir", + default="/home/ubuntu/Mooncake-efa/build-efa", + help="Path to the Mooncake build directory on both machines", + ) + parser.add_argument( + "--duration", + type=int, + default=10, + help="Benchmark duration in seconds per measurement point (default: 10)", + ) + parser.add_argument( + "--output", + default="efa_latency_bench.png", + help="Output chart filename (default: efa_latency_bench.png)", + ) + parser.add_argument( + "--cache_sizes", + default="1,5,10,20,30,40,50,60,70,80,90,100", + help="Comma-separated cache sizes in GB (default: 1,5,10,20,...,100)", + ) + parser.add_argument( + "--operation", + default="write", + choices=["read", "write"], + help="Transfer operation type (default: write)", + ) + parser.add_argument( + "--ssh_user", + default="ubuntu", + help="SSH username for remote connections (default: ubuntu)", + ) + parser.add_argument( + "--ssh_opts", + default="-o StrictHostKeyChecking=no -o ConnectTimeout=10", + help="Additional SSH options", + ) + parser.add_argument( + "--annotation_gb", + type=int, + default=40, + help="Cache size (GB) at which to annotate speedup ratio (default: 40)", + ) + parser.add_argument( + "--buffer_size", + type=int, + default=1073741824, + help="Buffer size in bytes for bench tool (default: 1073741824 = 1 GiB)", + ) + return parser.parse_args() + + +# Benchmark configurations +CONFIGS = [ + { + "label": "EFA (tuned)", + "env": {"MC_SLICE_SIZE": "262144"}, + "flags": { + "threads": 48, + "block_size": 131072, # 128KB + "batch_size": 128, + }, + "color": "#1f77b4", # blue + "marker": "o", + }, + { + "label": "EFA (default)", + "env": {}, + "flags": { + "threads": 8, + "block_size": 65536, # 64KB + "batch_size": 128, + }, + "color": "#ff7f0e", # orange + "marker": "^", + }, +] + + +def run_ssh(host, command, user="ubuntu", ssh_opts="", timeout=None): + """Run a command on a remote host via SSH.""" + ssh_args = [ + "ssh", + *ssh_opts.split(), + f"{user}@{host}", + command, + ] + try: + result = subprocess.run( + ssh_args, capture_output=True, text=True, timeout=timeout, + ) + return result.returncode, result.stdout, result.stderr + except subprocess.TimeoutExpired: + return -1, "", "Command timed out" + + +def run_cmd(cmd, timeout=None): + """Run a local shell command.""" + try: + result = subprocess.run( + cmd, shell=True, capture_output=True, text=True, timeout=timeout, + ) + return result.returncode, result.stdout, result.stderr + except subprocess.TimeoutExpired: + return -1, "", "Command timed out" + + +def kill_bench(target_host, user, ssh_opts): + """Kill any running transfer_engine_bench processes.""" + kill_cmd = ( + "ps aux | grep '[t]ransfer_engine_bench' | awk '{print $2}' " + "| xargs -r kill 2>/dev/null; sleep 1; echo done" + ) + run_ssh(target_host, kill_cmd, user, ssh_opts, timeout=15) + run_cmd( + "ps aux | grep '[t]ransfer_engine_bench' | awk '{print $2}' " + "| xargs -r kill 2>/dev/null", + timeout=5, + ) + time.sleep(2) + + +def start_target(target_host, build_dir, buffer_size, user, ssh_opts): + """Start the target process via SSH. Returns target address or None.""" + bench_bin = os.path.join( + build_dir, "mooncake-transfer-engine/example/transfer_engine_bench" + ) + log_file = "/tmp/efa_bench_target.log" + + target_cmd = ( + f"cd {build_dir} && " + f"env MC_METADATA_SERVER=P2PHANDSHAKE " + f"{bench_bin} " + f"--mode=target --protocol=efa --metadata_server=P2PHANDSHAKE " + f"--buffer_size={buffer_size} " + f"> {log_file} 2>&1" + ) + ssh_args = [ + "ssh", "-n", + *ssh_opts.split(), + f"{user}@{target_host}", + target_cmd, + ] + subprocess.Popen( + ssh_args, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + stdin=subprocess.DEVNULL, + ) + + for _ in range(20): + time.sleep(1) + rc, stdout, _ = run_ssh( + target_host, + f"grep 'listening on' {log_file} 2>/dev/null", + user, ssh_opts, timeout=10, + ) + if rc == 0 and "listening on" in stdout: + match = re.search(r"listening on (\S+:\d+)", stdout) + if match: + return match.group(1) + + _, log_out, _ = run_ssh( + target_host, f"tail -20 {log_file}", user, ssh_opts, timeout=10 + ) + print(f" Target log:\n{log_out}", file=sys.stderr) + return None + + +def run_initiator(initiator_host, build_dir, target_addr, config, + buffer_size, duration, operation, user, ssh_opts): + """Run the initiator benchmark. Returns throughput in GB/s or None.""" + bench_bin = os.path.join( + build_dir, "mooncake-transfer-engine/example/transfer_engine_bench" + ) + flags = config["flags"] + + env_parts = ["MC_METADATA_SERVER=P2PHANDSHAKE"] + for key, val in config.get("env", {}).items(): + env_parts.append(f"{key}={val}") + env_str = " ".join(env_parts) + + bench_cmd = ( + f"cd {build_dir} && " + f"{env_str} {bench_bin} " + f"--mode=initiator --protocol=efa --metadata_server=P2PHANDSHAKE " + f"--segment_id={target_addr} " + f"--operation={operation} " + f"--duration={duration} " + f"--threads={flags['threads']} " + f"--block_size={flags['block_size']} " + f"--batch_size={flags['batch_size']} " + f"--buffer_size={buffer_size} " + f"--report_unit=GB " + f"2>&1" + ) + + timeout = duration + 60 + rc, stdout, stderr = run_ssh( + initiator_host, bench_cmd, user, ssh_opts, timeout=timeout + ) + + combined = stdout + "\n" + stderr + match = re.search(r"throughput\s+([\d.]+)\s+GB/s", combined) + if match: + return float(match.group(1)) + + print(f" WARNING: Could not parse throughput", file=sys.stderr) + lines = combined.strip().split("\n") + for line in lines[-3:]: + print(f" {line}", file=sys.stderr) + return None + + +def plot_results(results, cache_sizes_gb, annotation_gb, output_path): + """Generate and save the latency vs cache size chart.""" + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + fig, ax = plt.subplots(1, 1, figsize=(7, 5)) + + for entry in results: + ax.plot( + cache_sizes_gb, + entry["latencies"], + color=entry["color"], + marker=entry["marker"], + markersize=5, + linewidth=1.5, + label=entry["label"], + ) + + ax.set_xlabel("Cache Size (GB)", fontsize=12) + ax.set_ylabel("Latency (s)", fontsize=12) + ax.set_xlim(0, max(cache_sizes_gb)) + ax.set_ylim(bottom=0) + ax.grid(True, alpha=0.3) + + # Speedup annotation + if len(results) >= 2 and annotation_gb in cache_sizes_gb: + idx = cache_sizes_gb.index(annotation_gb) + tuned_lat = results[0]["latencies"][idx] + default_lat = results[1]["latencies"][idx] + + if tuned_lat > 0: + speedup = default_lat / tuned_lat + ax.axvline(x=annotation_gb, color="gray", linestyle="--", alpha=0.5) + mid_y = (tuned_lat + default_lat) / 2 + ax.annotate( + "", + xy=(annotation_gb + 1, tuned_lat), + xytext=(annotation_gb + 1, default_lat), + arrowprops=dict(arrowstyle="<->", color="black", lw=1.5), + ) + ax.text( + annotation_gb + 3, mid_y, + f"{speedup:.1f}x", + fontsize=11, fontweight="bold", va="center", + ) + + ax.text( + 0.02, 0.98, "8 x 400 Gbps EFA NICs", + transform=ax.transAxes, fontsize=10, + verticalalignment="top", fontweight="bold", + ) + ax.legend(fontsize=10, loc="upper left", bbox_to_anchor=(0.0, 0.90)) + + plt.tight_layout() + plt.savefig(output_path, dpi=150, bbox_inches="tight") + print(f"\nChart saved to: {output_path}") + + +def main(): + args = parse_args() + cache_sizes_gb = [float(x) for x in args.cache_sizes.split(",")] + + print("=" * 60) + print("EFA Latency vs Cache Size Benchmark") + print("=" * 60) + print(f" Target host: {args.target_host}") + print(f" Initiator host: {args.initiator_host}") + print(f" Build dir: {args.build_dir}") + print(f" Duration: {args.duration}s per measurement") + print(f" Operation: {args.operation}") + print(f" Buffer size: {args.buffer_size} bytes") + print(f" Cache sizes: {cache_sizes_gb} GB") + print(f" Output: {args.output}") + print(f" Configs: {len(CONFIGS)}") + n_runs = len(CONFIGS) * len(cache_sizes_gb) + print(f" Total runs: {n_runs} " + f"(each {args.duration}s + overhead)") + print() + + results = [] + + for ci, config in enumerate(CONFIGS): + label = config["label"] + print(f"[{ci+1}/{len(CONFIGS)}] Config: {label}") + print(f" threads={config['flags']['threads']}, " + f"block_size={config['flags']['block_size']}, " + f"batch_size={config['flags']['batch_size']}") + if config.get("env"): + print(f" env: {config['env']}") + print() + + throughputs = [] + latencies = [] + + # Start target once for all cache sizes of this config + kill_bench(args.target_host, args.ssh_user, args.ssh_opts) + print(f" Starting target...", end="", flush=True) + target_addr = start_target( + args.target_host, args.build_dir, args.buffer_size, + args.ssh_user, args.ssh_opts, + ) + if not target_addr: + print(" FAILED") + continue + print(f" ready ({target_addr})") + print() + + for si, cache_gb in enumerate(cache_sizes_gb): + tag = f" [{si+1}/{len(cache_sizes_gb)}] {cache_gb:6.0f} GB" + print(f"{tag} ...", end="", flush=True) + + tp = run_initiator( + args.initiator_host, args.build_dir, + target_addr, config, args.buffer_size, + args.duration, args.operation, + args.ssh_user, args.ssh_opts, + ) + + if tp is None: + print(" FAILED") + throughputs.append(None) + latencies.append(None) + continue + + lat = cache_gb / tp + throughputs.append(tp) + latencies.append(lat) + print(f" {tp:7.2f} GB/s lat={lat:.3f}s") + + # Cleanup + kill_bench(args.target_host, args.ssh_user, args.ssh_opts) + + valid = [t for t in throughputs if t is not None] + if not valid: + print(f" ERROR: No results for {label}\n", file=sys.stderr) + continue + + results.append({ + "label": label, + "throughputs": throughputs, + "latencies": latencies, + "color": config["color"], + "marker": config["marker"], + }) + print() + + if not results: + print("ERROR: No results collected.", file=sys.stderr) + sys.exit(1) + + # Summary table + print("=" * 60) + print("Results Summary") + print("=" * 60) + header = f"{'Cache(GB)':>10}" + for r in results: + header += f" {r['label']:>24}" + print(header) + print("-" * len(header)) + for i, cache_gb in enumerate(cache_sizes_gb): + row = f"{cache_gb:>10.0f}" + for r in results: + tp = r["throughputs"][i] + lat = r["latencies"][i] + if tp is not None: + row += f" {tp:7.2f} GB/s lat={lat:.3f}s" + else: + row += f" {'N/A':>24}" + print(row) + print() + + # Plot only valid points + plot_cache = [] + plot_data = [{**r, "latencies": []} for r in results] + for i, cache_gb in enumerate(cache_sizes_gb): + if all(r["latencies"][i] is not None for r in results): + plot_cache.append(cache_gb) + for j in range(len(results)): + plot_data[j]["latencies"].append(results[j]["latencies"][i]) + + plot_results(plot_data, plot_cache, args.annotation_gb, args.output) + + +if __name__ == "__main__": + main() diff --git a/mooncake-transfer-engine/example/transfer_engine_bench.cpp b/mooncake-transfer-engine/example/transfer_engine_bench.cpp index 6ce1eba6..46f588ff 100644 --- a/mooncake-transfer-engine/example/transfer_engine_bench.cpp +++ b/mooncake-transfer-engine/example/transfer_engine_bench.cpp @@ -88,7 +88,7 @@ DEFINE_string(mode, "initiator", DEFINE_string(operation, "read", "Operation type: read or write"); DEFINE_string(protocol, "rdma", - "Transfer protocol: rdma|barex|tcp|nvlink|nvlink_intra|hip"); + "Transfer protocol: rdma|barex|tcp|efa|nvlink|nvlink_intra|hip"); DEFINE_string(device_name, "mlx5_2", "Device name to use, valid if protocol=rdma"); @@ -458,6 +458,11 @@ static Transport *installTransportFromFlags(TransferEngine *engine) { args.get()[0] = const_cast(nic_priority_matrix.c_str()); args.get()[1] = nullptr; xport = engine->installTransport(FLAGS_protocol.c_str(), args.get()); + } else if (FLAGS_protocol == "efa") { + // EFA needs topology discovery to find devices, but auto_discovery + // would auto-install RDMA transport. Manually discover instead. + engine->getLocalTopology()->discover({}); + xport = engine->installTransport("efa", nullptr); } else if (FLAGS_protocol == "tcp" || FLAGS_protocol == "nvlink" || FLAGS_protocol == "hip" || FLAGS_protocol == "nvlink_intra" || FLAGS_protocol == "ubshmem") { diff --git a/mooncake-transfer-engine/include/transfer_metadata.h b/mooncake-transfer-engine/include/transfer_metadata.h index 3e21c6ff..dae07757 100644 --- a/mooncake-transfer-engine/include/transfer_metadata.h +++ b/mooncake-transfer-engine/include/transfer_metadata.h @@ -117,6 +117,9 @@ class TransferMetadata { #endif std::vector qp_num; std::string reply_msg; // on error +#ifdef USE_EFA + std::string efa_addr; // EFA endpoint address (hex encoded) +#endif }; struct NotifyDesc { diff --git a/mooncake-transfer-engine/include/transport/efa_transport/efa_context.h b/mooncake-transfer-engine/include/transport/efa_transport/efa_context.h new file mode 100644 index 00000000..1a1c056e --- /dev/null +++ b/mooncake-transfer-engine/include/transport/efa_transport/efa_context.h @@ -0,0 +1,177 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef EFA_CONTEXT_H +#define EFA_CONTEXT_H + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common.h" +#include "efa_transport.h" +#include "transport/transport.h" + +namespace mooncake { + +class EfaEndPoint; +class EfaTransport; + +struct EfaCq { + EfaCq() : cq(nullptr), outstanding(0) {} + struct fid_cq *cq; + volatile int outstanding; +}; + +struct EfaMemoryRegionMeta { + void *addr; + size_t length; + struct fid_mr *mr; + uint64_t key; +}; + +// Simple endpoint store for EFA +class EfaEndpointStore { + public: + std::shared_ptr get(const std::string &peer_nic_path); + // Atomically get-or-insert: returns existing endpoint or inserts new_ep. + // Prevents duplicate endpoint creation from concurrent callers. + std::shared_ptr getOrInsert( + const std::string &peer_nic_path, std::shared_ptr new_ep); + void add(const std::string &peer_nic_path, + std::shared_ptr endpoint); + void remove(const std::string &peer_nic_path); + int disconnectAll(); + size_t size() const; + + private: + mutable RWSpinlock lock_; + std::unordered_map> endpoints_; +}; + +// EfaContext represents the set of resources controlled by each local EFA +// device, including Memory Region, CQ, EndPoint, etc. using libfabric +class EfaContext { + public: + EfaContext(EfaTransport &engine, const std::string &device_name); + + ~EfaContext(); + + int construct(size_t num_cq_list = 1, size_t num_comp_channels = 1, + uint8_t port = 1, int gid_index = -1, size_t max_cqe = 4096, + int max_endpoints = 256); + + private: + int deconstruct(); + + public: + // Memory Region Management + int registerMemoryRegion(void *addr, size_t length, int access); + int unregisterMemoryRegion(void *addr); + int preTouchMemory(void *addr, size_t length); + uint64_t rkey(void *addr); + uint64_t lkey(void *addr); + void *mrDesc(void *addr); // Get MR descriptor for fi_write local_desc + + private: + int registerMemoryRegionInternal(void *addr, size_t length, int access, + EfaMemoryRegionMeta &mrMeta); + + public: + bool active() const { return active_; } + void set_active(bool flag) { active_ = flag; } + + public: + // EndPoint Management + std::shared_ptr endpoint(const std::string &peer_nic_path); + int deleteEndpoint(const std::string &peer_nic_path); + int disconnectAllEndpoints(); + size_t getTotalQPNumber() const; + + public: + // Access to engine for endpoint handshake + EfaTransport &engine() { return engine_; } + const EfaTransport &engine() const { return engine_; } + + // Submit slices for transfer + int submitPostSend(const std::vector &slice_list); + + // Poll completion queue for completed operations + int pollCq(int max_entries, int cq_index = 0); + + // Get CQ count + size_t cqCount() const { return cq_list_.size(); } + + // Get CQ outstanding count pointer + volatile int *cqOutstandingCount(int cq_index) { + if (cq_index < 0 || (size_t)cq_index >= cq_list_.size()) return nullptr; + return &cq_list_[cq_index]->outstanding; + } + + public: + // Device name, such as `rdmap0s2` + std::string deviceName() const { return device_name_; } + + // NIC Path, such as `192.168.3.76@rdmap0s2` + std::string nicPath() const; + + public: + // Libfabric accessors + struct fid_fabric *fabric() const { return fabric_; } + struct fid_domain *domain() const { return domain_; } + struct fid_av *av() const { return av_; } + struct fi_info *info() const { return fi_info_; } + std::string localAddr() const; + + // Compatibility methods (libfabric doesn't use lid/gid like ibverbs) + uint16_t lid() const { return 0; } + std::string gid() const { return localAddr(); } + + private: + EfaTransport &engine_; + std::string device_name_; + + // Libfabric objects + struct fi_info *fi_info_; + struct fi_info *hints_; + struct fid_fabric *fabric_; + struct fid_domain *domain_; + struct fid_av *av_; // Address vector for peer addressing + + bool active_; + + std::shared_ptr endpoint_store_; + std::vector> cq_list_; + + RWSpinlock mr_lock_; + std::unordered_map mr_map_; +}; + +} // namespace mooncake + +#endif // EFA_CONTEXT_H diff --git a/mooncake-transfer-engine/include/transport/efa_transport/efa_endpoint.h b/mooncake-transfer-engine/include/transport/efa_transport/efa_endpoint.h new file mode 100644 index 00000000..dd810e7a --- /dev/null +++ b/mooncake-transfer-engine/include/transport/efa_transport/efa_endpoint.h @@ -0,0 +1,164 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef EFA_ENDPOINT_H +#define EFA_ENDPOINT_H + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "common.h" +#include "efa_context.h" +#include "transfer_metadata.h" +#include "transport/transport.h" + +namespace mooncake { + +class EfaContext; + +// Custom context for libfabric operations - stores slice pointer for completion +// handling This struct MUST have fi_context as its first member +struct EfaOpContext { + struct fi_context fi_ctx; // Must be first member + Transport::Slice *slice; // Slice pointer for completion handling + volatile int *wr_depth; // Pointer to endpoint's wr_depth_ for CQ + // completion decrement +}; + +// EfaEndPoint represents a libfabric endpoint for EFA communication. +// Unlike RDMA QPs, EFA uses RDM (Reliable Datagram) endpoints with +// an address vector for peer addressing. +class EfaEndPoint { + public: + using HandShakeDesc = TransferMetadata::HandShakeDesc; + + enum Status { INITIALIZING, UNCONNECTED, CONNECTED }; + + EfaEndPoint(EfaContext &context); + ~EfaEndPoint(); + + // Construct endpoint with specified completion queue + int construct(struct fid_cq *cq, size_t num_qp_list = 1, size_t max_sge = 4, + size_t max_wr = 256, size_t max_inline = 64); + + private: + int deconstruct(); + + public: + void setPeerNicPath(const std::string &peer_nic_path); + + int setupConnectionsByActive(); + + int setupConnectionsByActive(const std::string &peer_nic_path) { + setPeerNicPath(peer_nic_path); + return setupConnectionsByActive(); + } + + int setupConnectionsByPassive(const HandShakeDesc &peer_desc, + HandShakeDesc &local_desc); + + bool hasOutstandingSlice() const; + + bool active() const { return active_; } + + void set_active(bool flag) { + RWSpinlock::WriteGuard guard(lock_); + active_ = flag; + if (!flag) inactive_time_ = getCurrentTimeInNano(); + } + + double inactiveTime() { + if (active_) return 0.0; + return (getCurrentTimeInNano() - inactive_time_) / 1000000000.0; + } + + public: + bool connected() const { + return status_.load(std::memory_order_relaxed) == CONNECTED; + } + + void disconnect(); + int destroyQP(); + + private: + void disconnectUnlocked(); + + public: + const std::string toString() const; + + // Submit RDMA write/read operations via libfabric + int submitPostSend(std::vector &slice_list, + std::vector &failed_slice_list); + + // Get the number of endpoints (always 1 for EFA RDM) + size_t getQPNumber() const { return 1; } + + // Get local endpoint address for handshake + std::string getLocalAddr() const; + + // Get peer's fi_addr + fi_addr_t getPeerFiAddr() const { return peer_fi_addr_; } + + EfaContext &context() { return context_; } + + private: + // Setup connection using peer's address from handshake + int doSetupConnection(const std::string &peer_addr, + std::string *reply_msg = nullptr); + + // Insert peer address into address vector + int insertPeerAddr(const std::string &peer_addr); + + private: + EfaContext &context_; + std::atomic status_; + + RWSpinlock lock_; + std::string peer_nic_path_; + + // Libfabric endpoint + struct fid_ep *ep_; + struct fid_cq *tx_cq_; + struct fid_cq *rx_cq_; + fi_addr_t peer_fi_addr_; // Peer's address in the AV + + // Local endpoint address (for handshake) + std::vector local_addr_; + size_t local_addr_len_; + + volatile int wr_depth_; + int max_wr_depth_; + volatile int *cq_outstanding_; + + // Spinlock to serialize fi_write calls on this endpoint. + // libfabric RDM endpoints are not thread-safe by default. + std::atomic_flag post_lock_ = ATOMIC_FLAG_INIT; + + volatile bool active_; + volatile uint64_t inactive_time_; +}; + +} // namespace mooncake + +#endif // EFA_ENDPOINT_H diff --git a/mooncake-transfer-engine/include/transport/efa_transport/efa_transport.h b/mooncake-transfer-engine/include/transport/efa_transport/efa_transport.h new file mode 100644 index 00000000..50294237 --- /dev/null +++ b/mooncake-transfer-engine/include/transport/efa_transport/efa_transport.h @@ -0,0 +1,149 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef EFA_TRANSPORT_H_ +#define EFA_TRANSPORT_H_ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "topology.h" +#include "transfer_metadata.h" +#include "transport/transport.h" + +namespace mooncake { + +class EfaContext; +class EfaEndPoint; +class TransferMetadata; + +class EfaTransport : public Transport { + friend class EfaContext; + friend class EfaEndPoint; + + public: + using BufferDesc = TransferMetadata::BufferDesc; + using SegmentDesc = TransferMetadata::SegmentDesc; + using HandShakeDesc = TransferMetadata::HandShakeDesc; + + public: + EfaTransport(); + + ~EfaTransport(); + + int install(std::string &local_server_name, + std::shared_ptr meta, + std::shared_ptr topo) override; + + const char *getName() const override { return "efa"; } + + int registerLocalMemory(void *addr, size_t length, + const std::string &location, bool remote_accessible, + bool update_metadata) override; + + int unregisterLocalMemory(void *addr, bool update_metadata = true) override; + + int registerLocalMemoryBatch(const std::vector &buffer_list, + const std::string &location) override; + + int unregisterLocalMemoryBatch( + const std::vector &addr_list) override; + + private: + // Internal version with force_sequential option to avoid nested parallelism + int registerLocalMemoryInternal(void *addr, size_t length, + const std::string &location, + bool remote_accessible, + bool update_metadata, + bool force_sequential); + + int unregisterLocalMemoryInternal(void *addr, bool update_metadata, + bool force_sequential); + + // TRANSFER + + Status submitTransfer(BatchID batch_id, + const std::vector &entries) override; + + Status submitTransferTask( + const std::vector &task_list) override; + + Status getTransferStatus(BatchID batch_id, + std::vector &status); + + Status getTransferStatus(BatchID batch_id, size_t task_id, + TransferStatus &status) override; + + SegmentID getSegmentID(const std::string &segment_name); + + private: + int allocateLocalSegmentID(); + + int preTouchMemory(void *addr, size_t length); + + public: + int onSetupEfaConnections(const HandShakeDesc &peer_desc, + HandShakeDesc &local_desc); + + int sendHandshake(const std::string &peer_server_name, + const HandShakeDesc &local_desc, + HandShakeDesc &peer_desc) { + return metadata_->sendHandshake(peer_server_name, local_desc, + peer_desc); + } + + const std::string &local_server_name() const { return local_server_name_; } + + std::shared_ptr meta() { return metadata_; } + + private: + int initializeEfaResources(); + + int startHandshakeDaemon(std::string &local_server_name); + + public: + static int selectDevice(SegmentDesc *desc, uint64_t offset, size_t length, + int &buffer_id, int &device_id, int retry_cnt = 0); + static int selectDevice(SegmentDesc *desc, uint64_t offset, size_t length, + std::string_view hint, int &buffer_id, + int &device_id, int retry_cnt = 0); + + private: + // Start/stop CQ polling worker threads + void startWorkerThreads(); + void stopWorkerThreads(); + void workerThreadFunc(int thread_id); + + private: + std::vector> context_list_; + std::shared_ptr local_topology_; + + // CQ polling worker threads + std::atomic worker_running_{false}; + std::vector worker_threads_; +}; + +} // namespace mooncake + +#endif // EFA_TRANSPORT_H_ diff --git a/mooncake-transfer-engine/src/CMakeLists.txt b/mooncake-transfer-engine/src/CMakeLists.txt index 627c9d56..f05401fc 100644 --- a/mooncake-transfer-engine/src/CMakeLists.txt +++ b/mooncake-transfer-engine/src/CMakeLists.txt @@ -2,6 +2,8 @@ file(GLOB ENGINE_SOURCES "*.cpp") add_subdirectory(common) add_subdirectory(transport) +# EFA library path is set globally via common.cmake (LIBFABRIC_LIB_DIR) + set(CMAKE_INSTALL_RPATH_USE_LINK_PATH TRUE) if(USE_HIP) @@ -105,3 +107,8 @@ if(USE_INTRA_NVLINK) message(STATUS "Enabled USE_INTRA_NVLINK support") target_compile_definitions(transfer_engine PUBLIC USE_INTRA_NVLINK) endif() + +if(USE_EFA) + message(STATUS "Enabled USE_EFA (AWS Elastic Fabric Adapter) support") + target_link_libraries(transfer_engine PUBLIC fabric efa_transport) +endif() diff --git a/mooncake-transfer-engine/src/multi_transport.cpp b/mooncake-transfer-engine/src/multi_transport.cpp index c9425192..f1445ded 100644 --- a/mooncake-transfer-engine/src/multi_transport.cpp +++ b/mooncake-transfer-engine/src/multi_transport.cpp @@ -52,6 +52,9 @@ #ifdef USE_UBSHMEM #include "transport/ascend_transport/ubshmem_transport/ubshmem_transport.h" #endif +#ifdef USE_EFA +#include "transport/efa_transport/efa_transport.h" +#endif #include @@ -283,6 +286,11 @@ Transport *MultiTransport::installTransport(const std::string &proto, transport = new UBShmemTransport(); } #endif +#ifdef USE_EFA + else if (std::string(proto) == "efa") { + transport = new EfaTransport(); + } +#endif if (!transport) { LOG(ERROR) << "Unsupported transport " << proto diff --git a/mooncake-transfer-engine/src/transfer_metadata.cpp b/mooncake-transfer-engine/src/transfer_metadata.cpp index 47ed804d..c6015006 100644 --- a/mooncake-transfer-engine/src/transfer_metadata.cpp +++ b/mooncake-transfer-engine/src/transfer_metadata.cpp @@ -62,6 +62,9 @@ struct TransferHandshakeUtil { for (const auto &qp : desc.qp_num) qpNums.append(qp); root["qp_num"] = qpNums; root["reply_msg"] = desc.reply_msg; +#ifdef USE_EFA + root["efa_addr"] = desc.efa_addr; // EFA endpoint address +#endif return root; } @@ -74,6 +77,9 @@ struct TransferHandshakeUtil { for (const auto &qp : root["qp_num"]) desc.qp_num.push_back(qp.asUInt()); desc.reply_msg = root["reply_msg"].asString(); +#ifdef USE_EFA + desc.efa_addr = root["efa_addr"].asString(); // EFA endpoint address +#endif return 0; } }; @@ -164,7 +170,8 @@ int TransferMetadata::encodeSegmentDesc(const SegmentDesc &desc, segmentJSON["timestamp"] = getCurrentDateTime(); if (segmentJSON["protocol"] == "rdma" || - segmentJSON["protocol"] == "barex") { + segmentJSON["protocol"] == "barex" || + segmentJSON["protocol"] == "efa") { Json::Value devicesJSON(Json::arrayValue); for (const auto &device : desc.devices) { Json::Value deviceJSON; @@ -325,7 +332,8 @@ TransferMetadata::decodeSegmentDesc(Json::Value &segmentJSON, if (segmentJSON.isMember("timestamp")) desc->timestamp = segmentJSON["timestamp"].asString(); - if (desc->protocol == "rdma" || desc->protocol == "barex") { + if (desc->protocol == "rdma" || desc->protocol == "barex" || + desc->protocol == "efa") { for (const auto &deviceJSON : segmentJSON["devices"]) { DeviceDesc device; device.name = deviceJSON["name"].asString(); diff --git a/mooncake-transfer-engine/src/transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/CMakeLists.txt index 91e9bb14..1b72a721 100644 --- a/mooncake-transfer-engine/src/transport/CMakeLists.txt +++ b/mooncake-transfer-engine/src/transport/CMakeLists.txt @@ -56,3 +56,9 @@ if (USE_UBSHMEM) add_subdirectory(ascend_transport) target_sources(transport PUBLIC $) endif() + +if (USE_EFA) + add_subdirectory(efa_transport) + target_sources(transport PUBLIC $) + target_link_libraries(transport PRIVATE fabric) +endif() diff --git a/mooncake-transfer-engine/src/transport/efa_transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/efa_transport/CMakeLists.txt new file mode 100644 index 00000000..ff32f6eb --- /dev/null +++ b/mooncake-transfer-engine/src/transport/efa_transport/CMakeLists.txt @@ -0,0 +1,14 @@ +file(GLOB EFA_SOURCES "*.cpp") +add_library(efa_transport OBJECT ${EFA_SOURCES}) + +# Link against libfabric (fabric) instead of ibverbs for AWS EFA +target_link_libraries(efa_transport PRIVATE fabric glog::glog) + +target_include_directories(efa_transport PRIVATE + ${CMAKE_SOURCE_DIR}/mooncake-transfer-engine/include + ${CMAKE_SOURCE_DIR}/mooncake-transfer-engine/include/transport/efa_transport + ${CMAKE_SOURCE_DIR}/mooncake-common/include + ${LIBFABRIC_INCLUDE_DIR} +) + +# libfabric library path is set globally via common.cmake (LIBFABRIC_LIB_DIR) diff --git a/mooncake-transfer-engine/src/transport/efa_transport/efa_context.cpp b/mooncake-transfer-engine/src/transport/efa_transport/efa_context.cpp new file mode 100644 index 00000000..20dc3f1e --- /dev/null +++ b/mooncake-transfer-engine/src/transport/efa_transport/efa_context.cpp @@ -0,0 +1,562 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "transport/efa_transport/efa_context.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "config.h" +#include "transport/efa_transport/efa_endpoint.h" +#include "transport/efa_transport/efa_transport.h" +#include "transport/transport.h" + +namespace mooncake { + +// EfaEndpointStore implementation +std::shared_ptr EfaEndpointStore::get( + const std::string &peer_nic_path) { + RWSpinlock::ReadGuard guard(lock_); + auto it = endpoints_.find(peer_nic_path); + if (it != endpoints_.end()) { + return it->second; + } + return nullptr; +} + +std::shared_ptr EfaEndpointStore::getOrInsert( + const std::string &peer_nic_path, std::shared_ptr new_ep) { + RWSpinlock::WriteGuard guard(lock_); + auto it = endpoints_.find(peer_nic_path); + if (it != endpoints_.end()) { + return it->second; // Another thread already created it + } + endpoints_[peer_nic_path] = new_ep; + return new_ep; +} + +void EfaEndpointStore::add(const std::string &peer_nic_path, + std::shared_ptr endpoint) { + RWSpinlock::WriteGuard guard(lock_); + endpoints_[peer_nic_path] = endpoint; +} + +void EfaEndpointStore::remove(const std::string &peer_nic_path) { + RWSpinlock::WriteGuard guard(lock_); + endpoints_.erase(peer_nic_path); +} + +int EfaEndpointStore::disconnectAll() { + RWSpinlock::WriteGuard guard(lock_); + for (auto &entry : endpoints_) { + if (entry.second) { + entry.second->disconnect(); + } + } + return 0; +} + +size_t EfaEndpointStore::size() const { + RWSpinlock::ReadGuard guard(lock_); + return endpoints_.size(); +} + +// EfaContext implementation +EfaContext::EfaContext(EfaTransport &engine, const std::string &device_name) + : engine_(engine), + device_name_(device_name), + fi_info_(nullptr), + hints_(nullptr), + fabric_(nullptr), + domain_(nullptr), + av_(nullptr), + active_(true) {} + +EfaContext::~EfaContext() { + if (fabric_) deconstruct(); +} + +int EfaContext::construct(size_t num_cq_list, size_t num_comp_channels, + uint8_t port, int gid_index, size_t max_cqe, + int max_endpoints) { + endpoint_store_ = std::make_shared(); + + // Setup hints for EFA provider + hints_ = fi_allocinfo(); + if (!hints_) { + LOG(ERROR) << "Failed to allocate fi_info hints"; + return ERR_CONTEXT; + } + + hints_->caps = + FI_MSG | FI_RMA | FI_READ | FI_WRITE | FI_REMOTE_READ | FI_REMOTE_WRITE; + hints_->mode = FI_CONTEXT; + hints_->ep_attr->type = FI_EP_RDM; // EFA uses RDM endpoints + hints_->fabric_attr->prov_name = strdup("efa"); + + // Specify the domain (device) name - append "-rdm" for RDM endpoint + std::string domain_name = device_name_ + "-rdm"; + hints_->domain_attr->name = strdup(domain_name.c_str()); + hints_->domain_attr->mr_mode = + FI_MR_LOCAL | FI_MR_VIRT_ADDR | FI_MR_ALLOCATED | FI_MR_PROV_KEY; + hints_->domain_attr->threading = FI_THREAD_SAFE; + + // Get fabric info + int ret = + fi_getinfo(FI_VERSION(1, 14), nullptr, nullptr, 0, hints_, &fi_info_); + if (ret) { + LOG(ERROR) << "fi_getinfo failed for device " << device_name_ << ": " + << fi_strerror(-ret); + fi_freeinfo(hints_); + hints_ = nullptr; + return ERR_CONTEXT; + } + + // Open fabric + ret = fi_fabric(fi_info_->fabric_attr, &fabric_, nullptr); + if (ret) { + LOG(ERROR) << "fi_fabric failed: " << fi_strerror(-ret); + fi_freeinfo(fi_info_); + fi_freeinfo(hints_); + fi_info_ = nullptr; + hints_ = nullptr; + return ERR_CONTEXT; + } + + // Open domain + ret = fi_domain(fabric_, fi_info_, &domain_, nullptr); + if (ret) { + LOG(ERROR) << "fi_domain failed: " << fi_strerror(-ret); + fi_close(&fabric_->fid); + fi_freeinfo(fi_info_); + fi_freeinfo(hints_); + fabric_ = nullptr; + fi_info_ = nullptr; + hints_ = nullptr; + return ERR_CONTEXT; + } + + // Create address vector + struct fi_av_attr av_attr = {}; + av_attr.type = FI_AV_TABLE; + av_attr.count = max_endpoints; + + ret = fi_av_open(domain_, &av_attr, &av_, nullptr); + if (ret) { + LOG(ERROR) << "fi_av_open failed: " << fi_strerror(-ret); + fi_close(&domain_->fid); + fi_close(&fabric_->fid); + fi_freeinfo(fi_info_); + fi_freeinfo(hints_); + domain_ = nullptr; + fabric_ = nullptr; + fi_info_ = nullptr; + hints_ = nullptr; + return ERR_CONTEXT; + } + + // Create completion queues + cq_list_.resize(num_cq_list); + for (size_t i = 0; i < num_cq_list; ++i) { + auto cq = std::make_shared(); + + struct fi_cq_attr cq_attr = {}; + cq_attr.size = max_cqe; + cq_attr.format = FI_CQ_FORMAT_DATA; + cq_attr.wait_obj = FI_WAIT_NONE; + + ret = fi_cq_open(domain_, &cq_attr, &cq->cq, nullptr); + if (ret) { + LOG(ERROR) << "fi_cq_open failed: " << fi_strerror(-ret); + return ERR_CONTEXT; + } + cq_list_[i] = cq; + } + + LOG(INFO) << "EFA device (libfabric): " << device_name_ + << ", domain: " << fi_info_->domain_attr->name + << ", provider: " << fi_info_->fabric_attr->prov_name; + + return 0; +} + +int EfaContext::deconstruct() { + // Destroy all endpoints before closing domain/fabric/AV. + // Endpoints hold fi_ep handles that reference the domain, so they must + // be closed first. + endpoint_store_.reset(); + + { + RWSpinlock::WriteGuard guard(mr_lock_); + for (auto &entry : mr_map_) { + if (entry.second.mr) { + fi_close(&entry.second.mr->fid); + } + } + mr_map_.clear(); + } + + for (auto &cq : cq_list_) { + if (cq && cq->cq) { + fi_close(&cq->cq->fid); + cq->cq = nullptr; + } + } + cq_list_.clear(); + + if (av_) { + fi_close(&av_->fid); + av_ = nullptr; + } + + if (domain_) { + fi_close(&domain_->fid); + domain_ = nullptr; + } + + if (fabric_) { + fi_close(&fabric_->fid); + fabric_ = nullptr; + } + + if (fi_info_) { + fi_freeinfo(fi_info_); + fi_info_ = nullptr; + } + + if (hints_) { + fi_freeinfo(hints_); + hints_ = nullptr; + } + + return 0; +} + +int EfaContext::registerMemoryRegionInternal(void *addr, size_t length, + int access, + EfaMemoryRegionMeta &mrMeta) { + if (length > (size_t)globalConfig().max_mr_size) { + PLOG(WARNING) << "The buffer length exceeds device max_mr_size, " + << "shrink it to " << globalConfig().max_mr_size; + length = (size_t)globalConfig().max_mr_size; + } + + mrMeta.addr = addr; + mrMeta.length = length; + + // Convert access flags to libfabric flags + uint64_t fi_access = 0; + if (access & FI_READ) fi_access |= FI_READ; + if (access & FI_WRITE) fi_access |= FI_WRITE; + if (access & FI_REMOTE_READ) fi_access |= FI_REMOTE_READ; + if (access & FI_REMOTE_WRITE) fi_access |= FI_REMOTE_WRITE; + + // For EFA, we need local read/write and remote read/write + fi_access = FI_READ | FI_WRITE | FI_REMOTE_READ | FI_REMOTE_WRITE; + + int ret = fi_mr_reg(domain_, addr, length, fi_access, 0, 0, 0, &mrMeta.mr, + nullptr); + if (ret) { + LOG(ERROR) << "fi_mr_reg failed for " << addr << ": " + << fi_strerror(-ret); + return ERR_CONTEXT; + } + + mrMeta.key = fi_mr_key(mrMeta.mr); + + return 0; +} + +int EfaContext::registerMemoryRegion(void *addr, size_t length, int access) { + EfaMemoryRegionMeta mrMeta; + int ret = registerMemoryRegionInternal(addr, length, access, mrMeta); + if (ret != 0) { + return ret; + } + RWSpinlock::WriteGuard guard(mr_lock_); + mr_map_[(uint64_t)addr] = mrMeta; + return 0; +} + +int EfaContext::unregisterMemoryRegion(void *addr) { + RWSpinlock::WriteGuard guard(mr_lock_); + auto it = mr_map_.find((uint64_t)addr); + if (it != mr_map_.end()) { + if (it->second.mr) { + int ret = fi_close(&it->second.mr->fid); + if (ret) { + LOG(ERROR) << "Failed to unregister memory " << addr << ": " + << fi_strerror(-ret); + return ERR_CONTEXT; + } + } + mr_map_.erase(it); + } + return 0; +} + +int EfaContext::preTouchMemory(void *addr, size_t length) { + volatile char *ptr = (volatile char *)addr; + for (size_t i = 0; i < length; i += 4096) { + ptr[i] = ptr[i]; + } + return 0; +} + +uint64_t EfaContext::rkey(void *addr) { + RWSpinlock::ReadGuard guard(mr_lock_); + auto it = mr_map_.find((uint64_t)addr); + if (it != mr_map_.end() && it->second.mr) { + return it->second.key; + } + return 0; +} + +uint64_t EfaContext::lkey(void *addr) { + RWSpinlock::ReadGuard guard(mr_lock_); + auto it = mr_map_.find((uint64_t)addr); + if (it != mr_map_.end() && it->second.mr) { + return fi_mr_key(it->second.mr); + } + return 0; +} + +void *EfaContext::mrDesc(void *addr) { + RWSpinlock::ReadGuard guard(mr_lock_); + // Find the MR that contains this address + for (auto &entry : mr_map_) { + if ((uint64_t)addr >= entry.first && + (uint64_t)addr < entry.first + entry.second.length) { + if (entry.second.mr) { + return fi_mr_desc(entry.second.mr); + } + } + } + return nullptr; +} + +std::shared_ptr EfaContext::endpoint( + const std::string &peer_nic_path) { + if (!endpoint_store_) return nullptr; + + // Fast path: endpoint already exists + auto ep = endpoint_store_->get(peer_nic_path); + if (ep) return ep; + + // Slow path: create new endpoint, then atomically insert (or get existing + // if another thread raced us). getOrInsert prevents duplicate endpoints + // and duplicate AV entries for the same peer. + auto new_endpoint = std::make_shared(*this); + if (!cq_list_.empty() && cq_list_[0]) { + int ret = new_endpoint->construct(cq_list_[0]->cq); + if (ret != 0) { + LOG(ERROR) << "Failed to construct EFA endpoint"; + return nullptr; + } + } + new_endpoint->setPeerNicPath(peer_nic_path); + ep = endpoint_store_->getOrInsert(peer_nic_path, new_endpoint); + // If another thread won the race, new_endpoint is discarded (RAII cleanup) + return ep; +} + +int EfaContext::deleteEndpoint(const std::string &peer_nic_path) { + if (endpoint_store_) { + endpoint_store_->remove(peer_nic_path); + } + return 0; +} + +int EfaContext::disconnectAllEndpoints() { + if (endpoint_store_) { + return endpoint_store_->disconnectAll(); + } + return 0; +} + +size_t EfaContext::getTotalQPNumber() const { + return endpoint_store_ ? endpoint_store_->size() : 0; +} + +std::string EfaContext::nicPath() const { + return engine_.local_server_name() + "@" + device_name_; +} + +std::string EfaContext::localAddr() const { + // Return a hex string representation of the local address info + if (!fi_info_ || !fi_info_->src_addr) { + return ""; + } + + std::ostringstream oss; + const uint8_t *addr = static_cast(fi_info_->src_addr); + for (size_t i = 0; i < fi_info_->src_addrlen; ++i) { + oss << std::hex << std::setw(2) << std::setfill('0') << (int)addr[i]; + } + return oss.str(); +} + +int EfaContext::submitPostSend( + const std::vector &slice_list) { + // Route slices to appropriate endpoints for sending + // Group slices by peer NIC path + std::unordered_map> + slices_by_peer; + std::vector failed_slices; + + for (auto *slice : slice_list) { + if (!slice) continue; + + // Get peer segment descriptor to find dest_rkey and peer device info + auto peer_segment_desc = + engine_.meta()->getSegmentDescByID(slice->target_id); + if (!peer_segment_desc) { + LOG(ERROR) << "Cannot get segment descriptor for target " + << slice->target_id; + slice->markFailed(); + continue; + } + + // Find the buffer and device for this destination address + int buffer_id = -1, device_id = -1; + if (EfaTransport::selectDevice(peer_segment_desc.get(), + slice->rdma.dest_addr, slice->length, + buffer_id, device_id)) { + LOG(ERROR) << "Cannot select device for dest_addr " + << (void *)slice->rdma.dest_addr; + slice->markFailed(); + continue; + } + + // Set the remote key from the peer's registered memory region + slice->rdma.dest_rkey = + peer_segment_desc->buffers[buffer_id].rkey[device_id]; + + // Construct peer NIC path: "server_name@device_name" + std::string peer_nic_path = peer_segment_desc->name + "@" + + peer_segment_desc->devices[device_id].name; + slice->peer_nic_path = peer_nic_path; + + slices_by_peer[peer_nic_path].push_back(slice); + } + + // Now send to each peer endpoint + for (auto &entry : slices_by_peer) { + const std::string &peer_nic_path = entry.first; + auto &peer_slices = entry.second; + + // Get or create endpoint for this peer + auto ep = endpoint(peer_nic_path); + if (!ep) { + LOG(ERROR) << "Cannot create endpoint for peer " << peer_nic_path; + for (auto *slice : peer_slices) { + slice->markFailed(); + } + continue; + } + + // Submit to endpoint + std::vector failed_slice_list; + ep->submitPostSend(peer_slices, failed_slice_list); + + // Handle any slices that failed to post + for (auto *slice : failed_slice_list) { + slice->markFailed(); + } + } + + return 0; +} + +int EfaContext::pollCq(int max_entries, int cq_index) { + if (cq_index < 0 || (size_t)cq_index >= cq_list_.size()) { + return 0; + } + + struct fid_cq *cq = cq_list_[cq_index]->cq; + if (!cq) return 0; + + // Use fi_cq_data format for completions + struct fi_cq_data_entry entries[64]; + int to_poll = std::min(max_entries, 64); + + ssize_t ret = fi_cq_read(cq, entries, to_poll); + + if (ret > 0) { + // Process completions outside the lock (markSuccess / delete are safe) + std::unordered_map wr_depth_set; + for (ssize_t i = 0; i < ret; i++) { + EfaOpContext *op_ctx = + reinterpret_cast(entries[i].op_context); + if (op_ctx && op_ctx->slice) { + op_ctx->slice->markSuccess(); + if (op_ctx->wr_depth) { + wr_depth_set[op_ctx->wr_depth]++; + } + delete op_ctx; + } + } + for (auto &entry : wr_depth_set) { + __sync_fetch_and_sub(entry.first, entry.second); + } + __sync_fetch_and_sub(&cq_list_[cq_index]->outstanding, + static_cast(ret)); + return static_cast(ret); + } else if (ret == -FI_EAGAIN) { + return 0; + } else if (ret < 0) { + // CQ error - drain all queued error entries under the domain lock + int err_count = 0; + struct fi_cq_err_entry err_entry; + std::unordered_map wr_depth_set; + + while ((ret = fi_cq_readerr(cq, &err_entry, 0)) > 0) { + EfaOpContext *op_ctx = + reinterpret_cast(err_entry.op_context); + if (op_ctx && op_ctx->slice) { + LOG(ERROR) << "EFA CQ error: " + << fi_cq_strerror(cq, err_entry.prov_errno, + err_entry.err_data, nullptr, 0) + << " for slice at " << op_ctx->slice->source_addr; + op_ctx->slice->markFailed(); + if (op_ctx->wr_depth) { + wr_depth_set[op_ctx->wr_depth]++; + } + delete op_ctx; + } + err_count++; + } + + for (auto &entry : wr_depth_set) { + __sync_fetch_and_sub(entry.first, entry.second); + } + if (err_count > 0) { + __sync_fetch_and_sub(&cq_list_[cq_index]->outstanding, err_count); + } + return err_count; + } + + return 0; +} + +} // namespace mooncake diff --git a/mooncake-transfer-engine/src/transport/efa_transport/efa_endpoint.cpp b/mooncake-transfer-engine/src/transport/efa_transport/efa_endpoint.cpp new file mode 100644 index 00000000..e87fcb58 --- /dev/null +++ b/mooncake-transfer-engine/src/transport/efa_transport/efa_endpoint.cpp @@ -0,0 +1,441 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "transport/efa_transport/efa_endpoint.h" + +#include + +#include +#include +#include +#include +#include +#include + +#include "config.h" + +namespace mooncake { + +EfaEndPoint::EfaEndPoint(EfaContext &context) + : context_(context), + status_(INITIALIZING), + ep_(nullptr), + tx_cq_(nullptr), + rx_cq_(nullptr), + peer_fi_addr_(FI_ADDR_UNSPEC), + local_addr_len_(0), + wr_depth_(0), + max_wr_depth_(0), + cq_outstanding_(nullptr), + active_(true), + inactive_time_(0) {} + +EfaEndPoint::~EfaEndPoint() { + if (ep_) deconstruct(); +} + +int EfaEndPoint::construct(struct fid_cq *cq, size_t num_qp_list, + size_t max_sge, size_t max_wr, size_t max_inline) { + if (status_.load(std::memory_order_relaxed) != INITIALIZING) { + LOG(ERROR) << "EFA Endpoint has already been constructed"; + return ERR_ENDPOINT; + } + + tx_cq_ = cq; + rx_cq_ = cq; // Use same CQ for TX and RX + max_wr_depth_ = max_wr; + cq_outstanding_ = context_.cqOutstandingCount(0); + + // Create endpoint + int ret = fi_endpoint(context_.domain(), context_.info(), &ep_, nullptr); + if (ret) { + LOG(ERROR) << "fi_endpoint failed: " << fi_strerror(-ret); + return ERR_ENDPOINT; + } + + // Bind endpoint to AV + ret = fi_ep_bind(ep_, &context_.av()->fid, 0); + if (ret) { + LOG(ERROR) << "fi_ep_bind (av) failed: " << fi_strerror(-ret); + fi_close(&ep_->fid); + ep_ = nullptr; + return ERR_ENDPOINT; + } + + // Bind endpoint to TX CQ + ret = fi_ep_bind(ep_, &tx_cq_->fid, FI_TRANSMIT); + if (ret) { + LOG(ERROR) << "fi_ep_bind (tx_cq) failed: " << fi_strerror(-ret); + fi_close(&ep_->fid); + ep_ = nullptr; + return ERR_ENDPOINT; + } + + // Bind endpoint to RX CQ + ret = fi_ep_bind(ep_, &rx_cq_->fid, FI_RECV); + if (ret) { + LOG(ERROR) << "fi_ep_bind (rx_cq) failed: " << fi_strerror(-ret); + fi_close(&ep_->fid); + ep_ = nullptr; + return ERR_ENDPOINT; + } + + // Enable endpoint + ret = fi_enable(ep_); + if (ret) { + LOG(ERROR) << "fi_enable failed: " << fi_strerror(-ret); + fi_close(&ep_->fid); + ep_ = nullptr; + return ERR_ENDPOINT; + } + + // Get local endpoint address + local_addr_len_ = 64; // EFA addresses are typically 32 bytes + local_addr_.resize(local_addr_len_); + ret = fi_getname(&ep_->fid, local_addr_.data(), &local_addr_len_); + if (ret) { + LOG(ERROR) << "fi_getname failed: " << fi_strerror(-ret); + fi_close(&ep_->fid); + ep_ = nullptr; + return ERR_ENDPOINT; + } + local_addr_.resize(local_addr_len_); + + status_.store(UNCONNECTED, std::memory_order_relaxed); + return 0; +} + +int EfaEndPoint::deconstruct() { + if (ep_) { + fi_close(&ep_->fid); + ep_ = nullptr; + } + return 0; +} + +int EfaEndPoint::destroyQP() { return deconstruct(); } + +void EfaEndPoint::setPeerNicPath(const std::string &peer_nic_path) { + RWSpinlock::WriteGuard guard(lock_); + if (connected()) { + LOG(WARNING) << "Previous EFA connection will be discarded"; + disconnectUnlocked(); + } + peer_nic_path_ = peer_nic_path; +} + +std::string EfaEndPoint::getLocalAddr() const { + std::ostringstream oss; + for (size_t i = 0; i < local_addr_.size(); ++i) { + oss << std::hex << std::setw(2) << std::setfill('0') + << (int)local_addr_[i]; + } + return oss.str(); +} + +int EfaEndPoint::insertPeerAddr(const std::string &peer_addr) { + // Convert hex string to binary address + std::vector addr_bin; + addr_bin.reserve(peer_addr.size() / 2); + + for (size_t i = 0; i < peer_addr.size(); i += 2) { + std::string byte_str = peer_addr.substr(i, 2); + uint8_t byte = (uint8_t)strtol(byte_str.c_str(), nullptr, 16); + addr_bin.push_back(byte); + } + + // Insert into address vector + int ret = fi_av_insert(context_.av(), addr_bin.data(), 1, &peer_fi_addr_, 0, + nullptr); + if (ret != 1) { + LOG(ERROR) << "fi_av_insert failed: " << fi_strerror(-ret); + return ERR_ENDPOINT; + } + + return 0; +} + +int EfaEndPoint::setupConnectionsByActive() { + RWSpinlock::WriteGuard guard(lock_); + if (connected()) { + LOG(INFO) << "EFA Connection has been established"; + return 0; + } + + // Loopback mode + if (context_.nicPath() == peer_nic_path_) { + // For loopback, insert our own address + int ret = insertPeerAddr(getLocalAddr()); + if (ret != 0) { + return ret; + } + status_.store(CONNECTED, std::memory_order_release); + LOG(INFO) << "EFA loopback connection established: " << toString(); + return 0; + } + + // Exchange addresses via handshake + TransferMetadata::HandShakeDesc local_desc, peer_desc; + local_desc.local_nic_path = context_.nicPath(); + local_desc.peer_nic_path = peer_nic_path_; + // Store our EFA endpoint address in efa_addr field (hex encoded) + local_desc.efa_addr = getLocalAddr(); + + auto peer_server_name = getServerNameFromNicPath(peer_nic_path_); + auto peer_nic_name = getNicNameFromNicPath(peer_nic_path_); + if (peer_server_name.empty() || peer_nic_name.empty()) { + LOG(ERROR) << "Parse peer EFA nic path failed: " << peer_nic_path_; + return ERR_INVALID_ARGUMENT; + } + + int rc = context_.engine().sendHandshake(peer_server_name, local_desc, + peer_desc); + if (rc) return rc; + + if (peer_desc.efa_addr.empty()) { + LOG(ERROR) << "Peer did not provide EFA address in handshake"; + return ERR_REJECT_HANDSHAKE; + } + + // Insert peer's address into our AV + rc = insertPeerAddr(peer_desc.efa_addr); + if (rc != 0) { + return rc; + } + + status_.store(CONNECTED, std::memory_order_release); + VLOG(1) << "EFA connection established: " << toString() + << " peer_fi_addr=" << peer_fi_addr_; + return 0; +} + +int EfaEndPoint::setupConnectionsByPassive(const HandShakeDesc &peer_desc, + HandShakeDesc &local_desc) { + RWSpinlock::WriteGuard guard(lock_); + if (connected()) { + LOG(WARNING) << "Re-establish EFA connection: " << toString(); + disconnectUnlocked(); + } + + if (peer_desc.peer_nic_path != context_.nicPath() || + peer_desc.local_nic_path != peer_nic_path_) { + local_desc.reply_msg = "EFA nic path inconsistency"; + LOG(ERROR) << "Invalid argument: peer EFA nic path inconsistency" + << " peer_nic_path=" << peer_desc.peer_nic_path + << " context_.nicPath()=" << context_.nicPath() + << " local_nic_path=" << peer_desc.local_nic_path + << " peer_nic_path_=" << peer_nic_path_; + return ERR_REJECT_HANDSHAKE; + } + + // Insert peer's address from handshake + if (peer_desc.efa_addr.empty()) { + local_desc.reply_msg = "No EFA address provided"; + LOG(ERROR) << "Peer did not provide EFA address"; + return ERR_REJECT_HANDSHAKE; + } + + int ret = insertPeerAddr(peer_desc.efa_addr); + if (ret != 0) { + local_desc.reply_msg = "Failed to insert peer address"; + return ret; + } + + // Provide our address to peer (using efa_addr field, not reply_msg) + local_desc.local_nic_path = context_.nicPath(); + local_desc.peer_nic_path = peer_nic_path_; + local_desc.efa_addr = getLocalAddr(); + // reply_msg should be empty on success + + status_.store(CONNECTED, std::memory_order_release); + VLOG(1) << "EFA connection established (passive): " << toString(); + return 0; +} + +void EfaEndPoint::disconnect() { + RWSpinlock::WriteGuard guard(lock_); + disconnectUnlocked(); +} + +void EfaEndPoint::disconnectUnlocked() { + // For EFA RDM endpoints, we don't need to reset QP state + // Just remove peer from AV if needed and mark as disconnected + peer_fi_addr_ = FI_ADDR_UNSPEC; + status_.store(UNCONNECTED, std::memory_order_release); +} + +const std::string EfaEndPoint::toString() const { + return "EfaEndPoint[" + context_.nicPath() + " <-> " + peer_nic_path_ + "]"; +} + +bool EfaEndPoint::hasOutstandingSlice() const { return wr_depth_ > 0; } + +int EfaEndPoint::doSetupConnection(const std::string &peer_addr, + std::string *reply_msg) { + int ret = insertPeerAddr(peer_addr); + if (ret != 0) { + if (reply_msg) *reply_msg = "Failed to insert peer address into AV"; + return ret; + } + + status_.store(CONNECTED, std::memory_order_release); + return 0; +} + +int EfaEndPoint::submitPostSend( + std::vector &slice_list, + std::vector &failed_slice_list) { + if (!connected()) { + // Try to establish connection first + int ret = setupConnectionsByActive(); + if (ret != 0) { + // Move all slices to failed list + for (auto *slice : slice_list) { + failed_slice_list.push_back(slice); + } + slice_list.clear(); + return ret; + } + } + + // Process slices - using fi_write for RDMA write operations. + // Use atomic reserve-before-post to prevent CQ overflow when multiple + // threads post to endpoints sharing the same CQ. The CQ has a fixed + // capacity (max_cqe); if more completions arrive than it can hold, the + // provider silently drops them and those slices never complete (hang). + const int kMaxBackoffYields = 100000; + const int cq_limit = static_cast(globalConfig().max_cqe); + + for (auto it = slice_list.begin(); it != slice_list.end();) { + // --- Atomically reserve CQ and WR capacity before posting --- + // This eliminates the TOCTOU race where multiple threads pass the + // capacity check simultaneously and collectively overflow the CQ. + int backoff = 0; + bool reserved = false; + while (!reserved) { + // Try to reserve one WR slot + int cur_wr = wr_depth_; + if (cur_wr >= max_wr_depth_) { + if (++backoff > kMaxBackoffYields) goto timeout; + std::this_thread::yield(); + continue; + } + if (!__sync_bool_compare_and_swap(&wr_depth_, cur_wr, cur_wr + 1)) { + continue; // CAS failed, retry immediately + } + // WR slot reserved. Now try to reserve CQ slot. + if (cq_outstanding_) { + int cur_cq = *cq_outstanding_; + while (cur_cq < cq_limit) { + if (__sync_bool_compare_and_swap(cq_outstanding_, cur_cq, + cur_cq + 1)) { + reserved = true; + break; + } + cur_cq = *cq_outstanding_; + } + if (!reserved) { + // CQ full - release WR reservation and back off + __sync_fetch_and_sub(&wr_depth_, 1); + if (++backoff > kMaxBackoffYields) goto timeout; + std::this_thread::yield(); + continue; + } + } else { + reserved = true; + } + } + + { + Transport::Slice *slice = *it; + + // Get memory region descriptor for the local buffer + void *local_desc = context_.mrDesc(slice->source_addr); + if (!local_desc) { + LOG(ERROR) << "No MR descriptor found for address " + << slice->source_addr; + // Release reservations + __sync_fetch_and_sub(&wr_depth_, 1); + if (cq_outstanding_) __sync_fetch_and_sub(cq_outstanding_, 1); + failed_slice_list.push_back(slice); + it = slice_list.erase(it); + continue; + } + + // Allocate operation context to track the slice for completion + // Note: This memory is freed after CQ completion in pollCq + EfaOpContext *op_ctx = new EfaOpContext(); + memset(op_ctx, 0, sizeof(EfaOpContext)); + op_ctx->slice = slice; + op_ctx->wr_depth = &wr_depth_; + + // Serialize fi_write per-endpoint: concurrent fi_write on the + // same RDM endpoint corrupts provider state. Cross-endpoint + // safety is handled by the FI_THREAD_SAFE hint. + while (post_lock_.test_and_set(std::memory_order_acquire)) { + } + ssize_t ret = fi_write(ep_, + (void *)slice->source_addr, // local buffer + slice->length, local_desc, peer_fi_addr_, + slice->rdma.dest_addr, // remote address + slice->rdma.dest_rkey, // remote key + &op_ctx->fi_ctx); // context for completion + post_lock_.clear(std::memory_order_release); + + if (ret == 0) { + // Successfully posted - do NOT mark success here! + // Success is marked only after CQ completion in pollCq. + // WR and CQ reservations are already accounted for. + slice->status = Transport::Slice::PENDING; + it = slice_list.erase(it); + } else if (ret == -FI_EAGAIN) { + // Provider queue full - release reservations and retry + delete op_ctx; + __sync_fetch_and_sub(&wr_depth_, 1); + if (cq_outstanding_) __sync_fetch_and_sub(cq_outstanding_, 1); + std::this_thread::yield(); + // Don't advance iterator - retry the same slice + } else { + // Hard error - release reservations + LOG(ERROR) << "fi_write failed: " << fi_strerror(-ret) + << " (source=" << slice->source_addr + << ", len=" << slice->length + << ", dest=" << (void *)slice->rdma.dest_addr + << ", rkey=" << slice->rdma.dest_rkey << ")"; + delete op_ctx; + __sync_fetch_and_sub(&wr_depth_, 1); + if (cq_outstanding_) __sync_fetch_and_sub(cq_outstanding_, 1); + failed_slice_list.push_back(slice); + it = slice_list.erase(it); + } + } + continue; + + timeout: + LOG(WARNING) << "EFA submitPostSend: timed out waiting for CQ drain" + << " (wr_depth=" << wr_depth_ << ", max=" << max_wr_depth_ + << ", cq_outstanding=" + << (cq_outstanding_ ? *cq_outstanding_ : -1) + << ", max_cqe=" << cq_limit << ")"; + for (; it != slice_list.end(); ++it) { + failed_slice_list.push_back(*it); + } + slice_list.clear(); + return 0; + } + + return 0; +} + +} // namespace mooncake diff --git a/mooncake-transfer-engine/src/transport/efa_transport/efa_transport.cpp b/mooncake-transfer-engine/src/transport/efa_transport/efa_transport.cpp new file mode 100644 index 00000000..a4a2e3b1 --- /dev/null +++ b/mooncake-transfer-engine/src/transport/efa_transport/efa_transport.cpp @@ -0,0 +1,732 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "transport/efa_transport/efa_transport.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include + +#include "common.h" +#include "config.h" +#include "memory_location.h" +#include "topology.h" +#include "transport/efa_transport/efa_context.h" +#include "transport/efa_transport/efa_endpoint.h" + +namespace mooncake { + +EfaTransport::EfaTransport() { + LOG(INFO) << "[EFA] AWS Elastic Fabric Adapter transport initialized"; +} + +EfaTransport::~EfaTransport() { + stopWorkerThreads(); + metadata_->removeSegmentDesc(local_server_name_); + batch_desc_set_.clear(); + context_list_.clear(); +} + +void EfaTransport::startWorkerThreads() { + if (worker_running_) return; + + worker_running_ = true; + // One poller thread per context for responsive CQ draining under load + size_t num_threads = context_list_.size(); + for (size_t i = 0; i < num_threads; i++) { + worker_threads_.emplace_back(&EfaTransport::workerThreadFunc, this, i); + } + LOG(INFO) << "EfaTransport: Started " << num_threads + << " CQ polling worker threads"; +} + +void EfaTransport::stopWorkerThreads() { + if (!worker_running_) return; + + worker_running_ = false; + for (auto &thread : worker_threads_) { + if (thread.joinable()) { + thread.join(); + } + } + worker_threads_.clear(); + LOG(INFO) << "EfaTransport: Stopped CQ polling worker threads"; +} + +void EfaTransport::workerThreadFunc(int thread_id) { + const int kPollBatchSize = 64; + + while (worker_running_) { + bool did_work = false; + + // Poll CQs from all contexts + for (size_t ctx_idx = thread_id; ctx_idx < context_list_.size(); + ctx_idx += worker_threads_.size()) { + auto &context = context_list_[ctx_idx]; + if (!context || !context->active()) continue; + + for (size_t cq_idx = 0; cq_idx < context->cqCount(); cq_idx++) { + int completed = context->pollCq(kPollBatchSize, cq_idx); + if (completed > 0) { + did_work = true; + } + } + } + + // If no work was done, yield CPU briefly + if (!did_work) { + std::this_thread::yield(); + } + } +} + +int EfaTransport::install(std::string &local_server_name, + std::shared_ptr meta, + std::shared_ptr topo) { + if (topo == nullptr) { + LOG(ERROR) << "EfaTransport: missing topology"; + return ERR_INVALID_ARGUMENT; + } + + metadata_ = meta; + local_server_name_ = local_server_name; + local_topology_ = topo; + + auto ret = initializeEfaResources(); + if (ret) { + LOG(ERROR) << "EfaTransport: cannot initialize EFA resources"; + return ret; + } + + ret = allocateLocalSegmentID(); + if (ret) { + LOG(ERROR) << "Transfer engine cannot be initialized: cannot " + "allocate local segment"; + return ret; + } + + ret = startHandshakeDaemon(local_server_name); + if (ret) { + LOG(ERROR) << "EfaTransport: cannot start handshake daemon"; + return ret; + } + + ret = metadata_->updateLocalSegmentDesc(); + if (ret) { + LOG(ERROR) << "EfaTransport: cannot publish segments"; + return ret; + } + + // Start CQ polling worker threads + startWorkerThreads(); + + return 0; +} + +int EfaTransport::preTouchMemory(void *addr, size_t length) { + if (context_list_.size() == 0) { + return 0; + } + + auto hwc = std::thread::hardware_concurrency(); + auto num_threads = hwc > 64 ? 16 : std::min(hwc, 8u); + if (length > (size_t)globalConfig().max_mr_size) { + length = (size_t)globalConfig().max_mr_size; + } + size_t block_size = length / num_threads; + if (block_size == 0) { + return 0; + } + + std::vector threads; + threads.reserve(num_threads); + std::vector thread_results(num_threads, 0); + + for (size_t thread_i = 0; thread_i < num_threads; ++thread_i) { + void *block_addr = static_cast(addr) + thread_i * block_size; + threads.emplace_back([this, thread_i, block_addr, block_size, + &thread_results]() { + int ret = context_list_[0]->preTouchMemory(block_addr, block_size); + thread_results[thread_i] = ret; + }); + } + + for (auto &thread : threads) { + thread.join(); + } + + for (size_t i = 0; i < num_threads; ++i) { + if (thread_results[i] != 0) { + return thread_results[i]; + } + } + + return 0; +} + +int EfaTransport::registerLocalMemory(void *addr, size_t length, + const std::string &name, + bool remote_accessible, + bool update_metadata) { + return registerLocalMemoryInternal(addr, length, name, remote_accessible, + update_metadata, false); +} + +int EfaTransport::registerLocalMemoryInternal(void *addr, size_t length, + const std::string &name, + bool remote_accessible, + bool update_metadata, + bool force_sequential) { + (void)remote_accessible; + BufferDesc buffer_desc; + const int kBaseAccessRights = IBV_ACCESS_LOCAL_WRITE | + IBV_ACCESS_REMOTE_WRITE | + IBV_ACCESS_REMOTE_READ; + + int access_rights = kBaseAccessRights; + + bool do_pre_touch = context_list_.size() > 0 && + std::thread::hardware_concurrency() >= 4 && + length >= (size_t)4 * 1024 * 1024 * 1024; + if (do_pre_touch) { + int ret = preTouchMemory(addr, length); + if (ret != 0) { + return ret; + } + } + + int use_parallel_reg = 0; + if (!force_sequential) { + use_parallel_reg = globalConfig().parallel_reg_mr; + if (use_parallel_reg == -1) { + use_parallel_reg = context_list_.size() > 1 && do_pre_touch; + } + } + + auto reg_start = std::chrono::steady_clock::now(); + + if (use_parallel_reg) { + std::vector reg_threads; + reg_threads.reserve(context_list_.size()); + std::vector ret_codes(context_list_.size(), 0); + const int ar = access_rights; + + for (size_t i = 0; i < context_list_.size(); ++i) { + reg_threads.emplace_back([this, &ret_codes, i, addr, length, ar]() { + ret_codes[i] = + context_list_[i]->registerMemoryRegion(addr, length, ar); + }); + } + + for (auto &thread : reg_threads) { + thread.join(); + } + + for (size_t i = 0; i < ret_codes.size(); ++i) { + if (ret_codes[i] != 0) { + LOG(ERROR) + << "Failed to register memory region with EFA context " + << i; + return ret_codes[i]; + } + } + } else { + for (size_t i = 0; i < context_list_.size(); ++i) { + int ret = context_list_[i]->registerMemoryRegion(addr, length, + access_rights); + if (ret) { + LOG(ERROR) + << "Failed to register memory region with EFA context " + << i; + return ret; + } + } + } + + auto reg_end = std::chrono::steady_clock::now(); + auto reg_duration_ms = + std::chrono::duration_cast(reg_end - + reg_start) + .count(); + + if (globalConfig().trace) { + LOG(INFO) << "EFA registerMemoryRegion: addr=" << addr + << ", length=" << length + << ", contexts=" << context_list_.size() + << ", parallel=" << (use_parallel_reg ? "true" : "false") + << ", duration=" << reg_duration_ms << "ms"; + } + + // Collect keys from all contexts + for (auto &context : context_list_) { + buffer_desc.lkey.push_back(context->lkey(addr)); + buffer_desc.rkey.push_back(context->rkey(addr)); + } + + if (name == kWildcardLocation) { + bool only_first_page = true; + const std::vector entries = + getMemoryLocation(addr, length, only_first_page); + if (entries.empty()) return -1; + buffer_desc.name = entries[0].location; + } else { + buffer_desc.name = name; + } + + buffer_desc.addr = (uint64_t)addr; + buffer_desc.length = length; + int rc = metadata_->addLocalMemoryBuffer(buffer_desc, update_metadata); + if (rc) return rc; + return 0; +} + +int EfaTransport::unregisterLocalMemory(void *addr, bool update_metadata) { + return unregisterLocalMemoryInternal(addr, update_metadata, false); +} + +int EfaTransport::unregisterLocalMemoryInternal(void *addr, + bool update_metadata, + bool force_sequential) { + int rc = metadata_->removeLocalMemoryBuffer(addr, update_metadata); + if (rc) return rc; + + int use_parallel_unreg = 0; + if (!force_sequential) { + use_parallel_unreg = globalConfig().parallel_reg_mr; + if (use_parallel_unreg == -1) { + use_parallel_unreg = context_list_.size() > 1; + } + } + + if (use_parallel_unreg) { + std::vector unreg_threads; + unreg_threads.reserve(context_list_.size()); + std::vector ret_codes(context_list_.size(), 0); + + for (size_t i = 0; i < context_list_.size(); ++i) { + unreg_threads.emplace_back([this, &ret_codes, i, addr]() { + ret_codes[i] = context_list_[i]->unregisterMemoryRegion(addr); + }); + } + + for (auto &thread : unreg_threads) { + thread.join(); + } + + for (size_t i = 0; i < ret_codes.size(); ++i) { + if (ret_codes[i] != 0) { + LOG(ERROR) + << "Failed to unregister memory region with EFA context " + << i; + return ret_codes[i]; + } + } + } else { + for (size_t i = 0; i < context_list_.size(); ++i) { + int ret = context_list_[i]->unregisterMemoryRegion(addr); + if (ret) { + LOG(ERROR) + << "Failed to unregister memory region with EFA context " + << i; + return ret; + } + } + } + + return 0; +} + +int EfaTransport::allocateLocalSegmentID() { + auto desc = std::make_shared(); + if (!desc) return ERR_MEMORY; + desc->name = local_server_name_; + desc->protocol = "efa"; + for (auto &entry : context_list_) { + TransferMetadata::DeviceDesc device_desc; + device_desc.name = entry->deviceName(); + device_desc.lid = entry->lid(); + device_desc.gid = entry->gid(); + desc->devices.push_back(device_desc); + } + desc->topology = *(local_topology_.get()); + metadata_->addLocalSegment(LOCAL_SEGMENT_ID, local_server_name_, + std::move(desc)); + return 0; +} + +int EfaTransport::registerLocalMemoryBatch( + const std::vector &buffer_list, + const std::string &location) { + std::vector> results; + for (auto &buffer : buffer_list) { + results.emplace_back( + std::async(std::launch::async, [this, buffer, location]() -> int { + return registerLocalMemoryInternal(buffer.addr, buffer.length, + location, true, false, true); + })); + } + + for (size_t i = 0; i < buffer_list.size(); ++i) { + if (results[i].get()) { + LOG(WARNING) << "EfaTransport: Failed to register memory: addr " + << buffer_list[i].addr << " length " + << buffer_list[i].length; + } + } + + return metadata_->updateLocalSegmentDesc(); +} + +int EfaTransport::unregisterLocalMemoryBatch( + const std::vector &addr_list) { + std::vector> results; + for (auto &addr : addr_list) { + results.emplace_back( + std::async(std::launch::async, [this, addr]() -> int { + return unregisterLocalMemoryInternal(addr, false, true); + })); + } + + for (size_t i = 0; i < addr_list.size(); ++i) { + if (results[i].get()) + LOG(WARNING) << "EfaTransport: Failed to unregister memory: addr " + << addr_list[i]; + } + + return metadata_->updateLocalSegmentDesc(); +} + +Status EfaTransport::submitTransfer( + BatchID batch_id, const std::vector &entries) { + auto &batch_desc = *((BatchDesc *)(batch_id)); + if (batch_desc.task_list.size() + entries.size() > batch_desc.batch_size) { + LOG(ERROR) << "EfaTransport: Exceed the limitation of current batch's " + "capacity"; + return Status::InvalidArgument( + "EfaTransport: Exceed the limitation of capacity, batch id: " + + std::to_string(batch_id)); + } + + size_t task_id = batch_desc.task_list.size(); + batch_desc.task_list.resize(task_id + entries.size()); + std::vector task_list; + for (auto &task : batch_desc.task_list) task_list.push_back(&task); + return submitTransferTask(task_list); +} + +Status EfaTransport::submitTransferTask( + const std::vector &task_list) { + std::unordered_map, std::vector> + slices_to_post; + auto local_segment_desc = metadata_->getSegmentDescByID(LOCAL_SEGMENT_ID); + assert(local_segment_desc.get()); + const size_t kBlockSize = globalConfig().slice_size; + const int kMaxRetryCount = globalConfig().retry_cnt; + const size_t kFragmentSize = globalConfig().fragment_limit; + const size_t kSubmitWatermark = + globalConfig().max_wr * globalConfig().num_qp_per_ep; + uint64_t nr_slices; + for (size_t index = 0; index < task_list.size(); ++index) { + assert(task_list[index]); + auto &task = *task_list[index]; + nr_slices = 0; + assert(task.request); + auto &request = *task.request; + + auto request_buffer_id = -1, request_device_id = -1; + if (selectDevice(local_segment_desc.get(), (uint64_t)request.source, + request.length, request_buffer_id, + request_device_id)) { + request_buffer_id = -1; + request_device_id = -1; + } + + for (uint64_t offset = 0; offset < request.length; + offset += kBlockSize) { + Slice *slice = getSliceCache().allocate(); + assert(slice); + if (!slice->from_cache) { + nr_slices++; + } + + bool merge_final_slice = + request.length - offset <= kBlockSize + kFragmentSize; + + slice->source_addr = (char *)request.source + offset; + slice->length = + merge_final_slice ? request.length - offset : kBlockSize; + slice->opcode = request.opcode; + slice->rdma.dest_addr = request.target_offset + offset; + slice->rdma.retry_cnt = request.advise_retry_cnt; + slice->rdma.max_retry_cnt = kMaxRetryCount; + slice->task = &task; + slice->target_id = request.target_id; + slice->status = Slice::PENDING; + slice->ts = 0; + task.slice_list.push_back(slice); + + int buffer_id = -1, device_id = -1, + retry_cnt = request.advise_retry_cnt; + bool found_device = false; + if (request_buffer_id >= 0 && request_device_id >= 0) { + found_device = true; + buffer_id = request_buffer_id; + device_id = request_device_id; + } + while (retry_cnt < kMaxRetryCount && !found_device) { + if (selectDevice(local_segment_desc.get(), + (uint64_t)slice->source_addr, slice->length, + buffer_id, device_id, retry_cnt++)) + continue; + assert(device_id >= 0 && + static_cast(device_id) < context_list_.size()); + auto &context = context_list_[device_id]; + assert(context.get()); + if (!context->active()) continue; + assert(buffer_id >= 0 && + static_cast(buffer_id) < + local_segment_desc->buffers.size()); + assert(local_segment_desc->buffers[buffer_id].lkey.size() == + context_list_.size()); + found_device = true; + break; + } + if (!found_device) { + auto source_addr = slice->source_addr; + for (auto &entry : slices_to_post) + for (auto s : entry.second) getSliceCache().deallocate(s); + LOG(ERROR) << "Memory region not registered by any active EFA " + "device(s): " + << source_addr; + return Status::AddressNotRegistered( + "Memory region not registered by any active EFA " + "device(s): " + + std::to_string(reinterpret_cast(source_addr))); + } else { + auto &context = context_list_[device_id]; + if (!context->active()) { + LOG(ERROR) + << "EFA Device " << device_id << " is not active"; + return Status::InvalidArgument("EFA Device " + + std::to_string(device_id) + + " is not active"); + } + slice->rdma.source_lkey = + local_segment_desc->buffers[buffer_id].lkey[device_id]; + slices_to_post[context].push_back(slice); + task.total_bytes += slice->length; + __sync_fetch_and_add(&task.slice_count, 1); + } + + if (nr_slices >= kSubmitWatermark) { + for (auto &entry : slices_to_post) + entry.first->submitPostSend(entry.second); + slices_to_post.clear(); + nr_slices = 0; + } + + if (merge_final_slice) { + break; + } + } + } + + for (auto &entry : slices_to_post) + if (!entry.second.empty()) entry.first->submitPostSend(entry.second); + return Status::OK(); +} + +Status EfaTransport::getTransferStatus(BatchID batch_id, + std::vector &status) { + auto &batch_desc = *((BatchDesc *)(batch_id)); + const size_t task_count = batch_desc.task_list.size(); + status.resize(task_count); + for (size_t task_id = 0; task_id < task_count; task_id++) { + auto &task = batch_desc.task_list[task_id]; + status[task_id].transferred_bytes = task.transferred_bytes; + uint64_t success_slice_count = task.success_slice_count; + uint64_t failed_slice_count = task.failed_slice_count; + if (success_slice_count + failed_slice_count == task.slice_count) { + if (failed_slice_count) + status[task_id].s = TransferStatusEnum::FAILED; + else + status[task_id].s = TransferStatusEnum::COMPLETED; + task.is_finished = true; + } else { + status[task_id].s = TransferStatusEnum::WAITING; + } + } + return Status::OK(); +} + +Status EfaTransport::getTransferStatus(BatchID batch_id, size_t task_id, + TransferStatus &status) { + auto &batch_desc = *((BatchDesc *)(batch_id)); + const size_t task_count = batch_desc.task_list.size(); + if (task_id >= task_count) { + return Status::InvalidArgument( + "EfaTransport::getTransportStatus invalid argument, batch id: " + + std::to_string(batch_id)); + } + auto &task = batch_desc.task_list[task_id]; + status.transferred_bytes = task.transferred_bytes; + uint64_t success_slice_count = task.success_slice_count; + uint64_t failed_slice_count = task.failed_slice_count; + if (success_slice_count + failed_slice_count == task.slice_count) { + if (failed_slice_count) + status.s = TransferStatusEnum::FAILED; + else + status.s = TransferStatusEnum::COMPLETED; + task.is_finished = true; + } else { + status.s = TransferStatusEnum::WAITING; + } + return Status::OK(); +} + +EfaTransport::SegmentID EfaTransport::getSegmentID( + const std::string &segment_name) { + return metadata_->getSegmentID(segment_name); +} + +int EfaTransport::onSetupEfaConnections(const HandShakeDesc &peer_desc, + HandShakeDesc &local_desc) { + auto local_nic_name = getNicNameFromNicPath(peer_desc.peer_nic_path); + if (local_nic_name.empty()) return ERR_INVALID_ARGUMENT; + + // Find context by device name instead of using hca_list index, since + // context_list_ only contains EFA devices and may have different + // indexing than the full hca_list. + std::shared_ptr context; + for (auto &entry : context_list_) { + if (entry->deviceName() == local_nic_name) { + context = entry; + break; + } + } + if (!context) return ERR_INVALID_ARGUMENT; + + auto endpoint = context->endpoint(peer_desc.local_nic_path); + if (!endpoint) return ERR_ENDPOINT; + return endpoint->setupConnectionsByPassive(peer_desc, local_desc); +} + +int EfaTransport::initializeEfaResources() { + auto hca_list = local_topology_->getHcaList(); + + // Filter for EFA devices (names typically start with "rdmap" on AWS) + std::vector efa_devices; + std::vector non_efa_devices; + for (auto &device_name : hca_list) { + if (device_name.find("rdmap") != std::string::npos || + device_name.find("efa") != std::string::npos) { + efa_devices.push_back(device_name); + } else { + non_efa_devices.push_back(device_name); + } + } + + if (efa_devices.empty()) { + LOG(WARNING) << "EfaTransport: No EFA devices found, falling back to " + "all devices"; + efa_devices = hca_list; + non_efa_devices.clear(); + } + + // Disable non-EFA devices (e.g. ibp* IB devices) in the topology so that + // topology device indices stay aligned with context_list_ indices. + // Without this, selectDevice() can return an index from the full topology + // (which includes non-EFA devices), causing out-of-bounds access on + // context_list_ which only contains EFA devices. + for (auto &device_name : non_efa_devices) { + local_topology_->disableDevice(device_name); + LOG(INFO) << "EfaTransport: Disabled non-EFA device " << device_name + << " in topology"; + } + + for (auto &device_name : efa_devices) { + auto context = std::make_shared(*this, device_name); + auto &config = globalConfig(); + int ret = context->construct(config.num_cq_per_ctx, + config.num_comp_channels_per_ctx, + config.port, config.gid_index, + config.max_cqe, config.max_ep_per_ctx); + if (ret) { + local_topology_->disableDevice(device_name); + LOG(WARNING) << "EfaTransport: Disable device " << device_name; + } else { + context_list_.push_back(context); + LOG(INFO) << "EfaTransport: Initialized EFA device " << device_name; + } + } + if (context_list_.empty()) { + LOG(ERROR) << "EfaTransport: No available EFA devices"; + return ERR_DEVICE_NOT_FOUND; + } + return 0; +} + +int EfaTransport::startHandshakeDaemon(std::string &local_server_name) { + return metadata_->startHandshakeDaemon( + std::bind(&EfaTransport::onSetupEfaConnections, this, + std::placeholders::_1, std::placeholders::_2), + metadata_->localRpcMeta().rpc_port, metadata_->localRpcMeta().sockfd); +} + +int EfaTransport::selectDevice(SegmentDesc *desc, uint64_t offset, + size_t length, std::string_view hint, + int &buffer_id, int &device_id, + int retry_count) { + if (desc == nullptr) return ERR_ADDRESS_NOT_REGISTERED; + const auto &buffers = desc->buffers; + for (buffer_id = 0; buffer_id < static_cast(buffers.size()); + ++buffer_id) { + const auto &buffer = buffers[buffer_id]; + + if (offset < buffer.addr || length > buffer.length || + offset - buffer.addr > buffer.length - length) { + continue; + } + + device_id = + hint.empty() + ? desc->topology.selectDevice(buffer.name, retry_count) + : desc->topology.selectDevice(buffer.name, hint, retry_count); + if (device_id >= 0) return 0; + device_id = hint.empty() ? desc->topology.selectDevice( + kWildcardLocation, retry_count) + : desc->topology.selectDevice( + kWildcardLocation, hint, retry_count); + if (device_id >= 0) return 0; + } + return ERR_ADDRESS_NOT_REGISTERED; +} + +int EfaTransport::selectDevice(SegmentDesc *desc, uint64_t offset, + size_t length, int &buffer_id, int &device_id, + int retry_count) { + return selectDevice(desc, offset, length, "", buffer_id, device_id, + retry_count); +} + +} // namespace mooncake diff --git a/mooncake-transfer-engine/tests/CMakeLists.txt b/mooncake-transfer-engine/tests/CMakeLists.txt index 12a607b3..ca98ea6c 100644 --- a/mooncake-transfer-engine/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tests/CMakeLists.txt @@ -54,6 +54,12 @@ if (USE_UBSHMEM) add_test(NAME ubshmem_transport_test COMMAND ubshmem_transport_test) endif() +if (USE_EFA) + add_executable(efa_transport_test ${WORKSPACE}/efa_transport_test.cpp) + target_link_libraries(efa_transport_test PUBLIC transfer_engine gtest gtest_main) + add_test(NAME efa_transport_test COMMAND efa_transport_test) +endif() + add_executable(transfer_metadata_test ${WORKSPACE}/transfer_metadata_test.cpp) target_link_libraries(transfer_metadata_test PUBLIC transfer_engine gtest gtest_main) add_test(NAME transfer_metadata_test COMMAND transfer_metadata_test) diff --git a/mooncake-transfer-engine/tests/efa_transport_test.cpp b/mooncake-transfer-engine/tests/efa_transport_test.cpp new file mode 100644 index 00000000..066421ac --- /dev/null +++ b/mooncake-transfer-engine/tests/efa_transport_test.cpp @@ -0,0 +1,338 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include + +#include +#include + +#include "transfer_engine.h" +#include "transport/transport.h" + +using namespace mooncake; + +namespace mooncake { + +static void *allocateMemoryPool(size_t size, int socket_id) { + return numa_alloc_onnode(size, socket_id); +} + +static void freeMemoryPool(void *addr, size_t size) { numa_free(addr, size); } + +// --------------------------------------------------------------------------- +// EFA Transport Test Fixture +// +// This test uses the P2PHANDSHAKE metadata backend and performs loopback +// transfers (local_server_name == segment_id), similar to the TCP transport +// tests. It requires EFA hardware to be present (fi_info -p efa must succeed). +// +// Environment variables: +// MC_METADATA_SERVER - metadata backend (default: P2PHANDSHAKE) +// MC_LOCAL_SERVER_NAME - local server name (default: 127.0.0.1:12345) +// --------------------------------------------------------------------------- +class EFATransportTest : public ::testing::Test { + protected: + void SetUp() override { + google::InitGoogleLogging("EFATransportTest"); + FLAGS_logtostderr = 1; + + const char *env = std::getenv("MC_METADATA_SERVER"); + metadata_server_ = env ? env : "P2PHANDSHAKE"; + LOG(INFO) << "metadata_server: " << metadata_server_; + + env = std::getenv("MC_LOCAL_SERVER_NAME"); + local_server_name_ = env ? env : "127.0.0.1:12345"; + LOG(INFO) << "local_server_name: " << local_server_name_; + } + + void TearDown() override { google::ShutdownGoogleLogging(); } + + // Helper: create engine, install EFA transport, register memory + struct EngineSetup { + std::unique_ptr engine; + Transport *xport; + void *addr; + size_t buffer_size; + SegmentID segment_id; + }; + + EngineSetup createEngine(size_t buffer_size = 1ull << 30) { + EngineSetup s; + s.buffer_size = buffer_size; + + s.engine = std::make_unique(false); + // Manually discover topology to populate EFA device list + // (same pattern as the Python binding in transfer_engine_py.cpp) + s.engine->getLocalTopology()->discover({}); + auto hp = parseHostNameWithPort(local_server_name_); + int rc = s.engine->init(metadata_server_, local_server_name_, + hp.first.c_str(), hp.second); + EXPECT_EQ(rc, 0) << "engine->init failed"; + + s.xport = s.engine->installTransport("efa", nullptr); + EXPECT_NE(s.xport, nullptr) << "installTransport(\"efa\") failed"; + + s.addr = allocateMemoryPool(buffer_size, 0); + EXPECT_NE(s.addr, nullptr) << "allocateMemoryPool failed"; + + rc = s.engine->registerLocalMemory(s.addr, buffer_size, "cpu:0"); + EXPECT_EQ(rc, 0) << "registerLocalMemory failed"; + + // Use actual RPC address (P2PHANDSHAKE picks a random port) + auto actual_addr = s.engine->getLocalIpAndPort(); + s.segment_id = s.engine->openSegment(actual_addr); + return s; + } + + void destroyEngine(EngineSetup &s) { + if (s.engine && s.addr) { + s.engine->unregisterLocalMemory(s.addr); + } + if (s.addr) { + freeMemoryPool(s.addr, s.buffer_size); + s.addr = nullptr; + } + } + + // Helper: submit a single transfer and poll until completion + bool submitAndWait(TransferEngine *engine, SegmentID segment_id, + void *source, uint64_t target_offset, size_t length, + TransferRequest::OpCode opcode) { + auto batch_id = engine->allocateBatchID(1); + + TransferRequest entry; + entry.opcode = opcode; + entry.length = length; + entry.source = (uint8_t *)source; + entry.target_id = segment_id; + entry.target_offset = target_offset; + + Status s = engine->submitTransfer(batch_id, {entry}); + if (!s.ok()) { + LOG(ERROR) << "submitTransfer failed: " << s.ToString(); + engine->freeBatchID(batch_id); + return false; + } + + // Poll for completion with timeout + const int kMaxPollIterations = 1000000; + TransferStatus status; + for (int i = 0; i < kMaxPollIterations; ++i) { + s = engine->getTransferStatus(batch_id, 0, status); + if (!s.ok()) { + LOG(ERROR) << "getTransferStatus failed: " << s.ToString(); + engine->freeBatchID(batch_id); + return false; + } + if (status.s == TransferStatusEnum::COMPLETED) { + engine->freeBatchID(batch_id); + return true; + } + if (status.s == TransferStatusEnum::FAILED) { + LOG(ERROR) << "Transfer FAILED"; + engine->freeBatchID(batch_id); + return false; + } + } + LOG(ERROR) << "Transfer timed out"; + engine->freeBatchID(batch_id); + return false; + } + + std::string metadata_server_; + std::string local_server_name_; +}; + +// Test 1: Verify EFA transport can be installed +TEST_F(EFATransportTest, InstallTransport) { + auto engine = std::make_unique(false); + engine->getLocalTopology()->discover({}); + auto hp = parseHostNameWithPort(local_server_name_); + int rc = engine->init(metadata_server_, local_server_name_, + hp.first.c_str(), hp.second); + ASSERT_EQ(rc, 0); + + Transport *xport = engine->installTransport("efa", nullptr); + ASSERT_NE(xport, nullptr) + << "EFA transport should be installable on EFA hardware"; +} + +// Test 2: Basic loopback write +TEST_F(EFATransportTest, LoopbackWrite) { + auto setup = createEngine(); + + auto segment_desc = + setup.engine->getMetadata()->getSegmentDescByID(setup.segment_id); + ASSERT_NE(segment_desc, nullptr); + uint64_t remote_base = (uint64_t)segment_desc->buffers[0].addr; + + const size_t kDataLength = 4096; + + // Fill source buffer with known data + memset(setup.addr, 0xAB, kDataLength); + + bool ok = submitAndWait(setup.engine.get(), setup.segment_id, setup.addr, + remote_base, kDataLength, TransferRequest::WRITE); + EXPECT_TRUE(ok) << "Loopback write should succeed"; + + destroyEngine(setup); +} + +// Test 3: Write then read, verify data integrity +TEST_F(EFATransportTest, WriteAndRead) { + auto setup = createEngine(); + + auto segment_desc = + setup.engine->getMetadata()->getSegmentDescByID(setup.segment_id); + ASSERT_NE(segment_desc, nullptr); + uint64_t remote_base = (uint64_t)segment_desc->buffers[0].addr; + + const size_t kDataLength = 4096000; + uint8_t *buf = (uint8_t *)setup.addr; + + // Fill first half with random data + for (size_t i = 0; i < kDataLength; ++i) buf[i] = 'a' + lrand48() % 26; + + // Write local -> remote (loopback) + bool ok = submitAndWait(setup.engine.get(), setup.segment_id, buf, + remote_base, kDataLength, TransferRequest::WRITE); + ASSERT_TRUE(ok) << "Write should succeed"; + + // Read remote -> local (into second half of buffer) + ok = submitAndWait(setup.engine.get(), setup.segment_id, buf + kDataLength, + remote_base, kDataLength, TransferRequest::READ); + ASSERT_TRUE(ok) << "Read should succeed"; + + // Verify data integrity + EXPECT_EQ(0, memcmp(buf, buf + kDataLength, kDataLength)) + << "Read-back data should match written data"; + + destroyEngine(setup); +} + +// Test 4: Multiple sequential writes in a batch +TEST_F(EFATransportTest, MultiWrite) { + auto setup = createEngine(); + + auto segment_desc = + setup.engine->getMetadata()->getSegmentDescByID(setup.segment_id); + ASSERT_NE(segment_desc, nullptr); + uint64_t remote_base = (uint64_t)segment_desc->buffers[0].addr; + + const size_t kDataLength = 65536; + const int kBatchSize = 16; + + auto batch_id = setup.engine->allocateBatchID(kBatchSize); + + std::vector requests; + for (int i = 0; i < kBatchSize; ++i) { + TransferRequest entry; + entry.opcode = TransferRequest::WRITE; + entry.length = kDataLength; + entry.source = (uint8_t *)setup.addr + i * kDataLength; + entry.target_id = setup.segment_id; + entry.target_offset = remote_base + i * kDataLength; + requests.push_back(entry); + } + + Status s = setup.engine->submitTransfer(batch_id, requests); + ASSERT_TRUE(s.ok()) << "submitTransfer failed: " << s.ToString(); + + // Poll all tasks until completion + for (int task_id = 0; task_id < kBatchSize; ++task_id) { + TransferStatus status; + const int kMaxPollIterations = 1000000; + for (int i = 0; i < kMaxPollIterations; ++i) { + s = setup.engine->getTransferStatus(batch_id, task_id, status); + ASSERT_TRUE(s.ok()); + if (status.s == TransferStatusEnum::COMPLETED) break; + ASSERT_NE(status.s, TransferStatusEnum::FAILED) + << "Task " << task_id << " failed"; + } + ASSERT_EQ(status.s, TransferStatusEnum::COMPLETED) + << "Task " << task_id << " did not complete"; + } + + s = setup.engine->freeBatchID(batch_id); + ASSERT_TRUE(s.ok()); + + destroyEngine(setup); +} + +// Test 5: Stress test - multiple batches to verify no CQ overflow +TEST_F(EFATransportTest, StressMultipleBatches) { + auto setup = createEngine(); + + auto segment_desc = + setup.engine->getMetadata()->getSegmentDescByID(setup.segment_id); + ASSERT_NE(segment_desc, nullptr); + uint64_t remote_base = (uint64_t)segment_desc->buffers[0].addr; + + const size_t kDataLength = 65536; + const int kBatchSize = 8; + const int kNumBatches = 20; + + for (int batch = 0; batch < kNumBatches; ++batch) { + auto batch_id = setup.engine->allocateBatchID(kBatchSize); + + std::vector requests; + for (int i = 0; i < kBatchSize; ++i) { + TransferRequest entry; + entry.opcode = TransferRequest::WRITE; + entry.length = kDataLength; + entry.source = + (uint8_t *)setup.addr + (i + batch * kBatchSize) * kDataLength; + entry.target_id = setup.segment_id; + entry.target_offset = + remote_base + (i + batch * kBatchSize) * kDataLength; + requests.push_back(entry); + } + + Status s = setup.engine->submitTransfer(batch_id, requests); + ASSERT_TRUE(s.ok()) + << "Batch " << batch << " submitTransfer failed: " << s.ToString(); + + // Wait for all tasks in batch + for (int task_id = 0; task_id < kBatchSize; ++task_id) { + TransferStatus status; + const int kMaxPollIterations = 1000000; + for (int i = 0; i < kMaxPollIterations; ++i) { + s = setup.engine->getTransferStatus(batch_id, task_id, status); + ASSERT_TRUE(s.ok()); + if (status.s == TransferStatusEnum::COMPLETED) break; + ASSERT_NE(status.s, TransferStatusEnum::FAILED) + << "Batch " << batch << " task " << task_id << " failed"; + } + ASSERT_EQ(status.s, TransferStatusEnum::COMPLETED) + << "Batch " << batch << " task " << task_id + << " did not complete"; + } + + s = setup.engine->freeBatchID(batch_id); + ASSERT_TRUE(s.ok()); + } + + destroyEngine(setup); +} + +} // namespace mooncake + +int main(int argc, char **argv) { + gflags::ParseCommandLineFlags(&argc, &argv, false); + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +}