[TE] Add AWS EFA transport using libfabric (#1509)
* [TE] Add AWS EFA transport using libfabric
Add EfaTransport as a new transport backend for AWS Elastic Fabric
Adapter (EFA) devices. EFA exposes RDMA-like NICs but does not support
the full ibverbs QP API, so this transport uses libfabric's FI_EP_RDM
(Reliable Datagram Message) endpoint type instead.
Architecture (per EFA device):
EfaTransport → EfaContext → EfaEndPoint
- EfaContext: owns fabric/domain/AV/CQ/MR resources
- EfaEndPoint: one RDM endpoint per peer, with address-vector addressing
- Dedicated CQ poller thread per device for responsive completion draining
Key design decisions:
- FI_THREAD_SAFE requested from provider; per-endpoint spinlock on
fi_write as safety net for concurrent submission threads
- Atomic CAS reservation of CQ and WR capacity before posting fi_write
to prevent CQ overflow under high concurrency
- CQ error path drains all queued errors (fi_cq_readerr loop) before
returning, per libfabric semantics
- Retry-with-backoff on CQ/WR full instead of immediate slice failure
- Thread-safe endpoint creation via atomic getOrInsert to prevent
duplicate endpoints for the same peer
- Handshake exchanges EFA endpoint addresses via dedicated efa_addr
field in HandShakeDesc
Build: cmake -DUSE_EFA=ON (requires libfabric from AWS EFA installer)
Tested on p6-b200.48xlarge (8 EFA devices, 8×400 Gbps): 59.72 GB/s
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [TE] Add EFA unit tests and bench tool support
Add efa_transport_test with 5 test cases:
- InstallTransport: verify EFA transport installation
- LoopbackWrite: basic loopback write operation
- WriteAndRead: write then read with data integrity check
- MultiWrite: batch write (16 requests)
- StressMultipleBatches: stress test (20 batches × 8 requests)
Add --protocol=efa support to transfer_engine_bench with manual
topology discovery (EFA needs explicit discover() since
TransferEngine(false) skips auto-discovery).
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Docs] Add EFA transport documentation
Add comprehensive EFA transport documentation covering:
- Prerequisites and build instructions
- Usage with vLLM (prefill/decode disaggregation)
- Unit test descriptions and environment variables
- Benchmark results on p6-b200.48xlarge: 59.72 GB/s (EFA) vs
9.5 GB/s (TCP iperf3) vs 0.11 GB/s (Mooncake TCP transport)
- EFA vs RoCE RDMA comparison table
- Thread safety design notes
- Troubleshooting guide
Add EfaTransport to the transfer-engine index toctree and supported
transport lists.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [TE] Address PR review: ifdef EFA fields, use find_package for libfabric
- Wrap efa_addr in HandShakeDesc with #ifdef USE_EFA in transfer_metadata.h
- Wrap efa_addr serialization/deserialization with #ifdef USE_EFA in transfer_metadata.cpp
- Replace hardcoded /opt/amazon/efa paths with find_path/find_library in common.cmake
- Remove redundant hardcoded EFA paths from all CMakeLists.txt files
- Fix git clone URL in efa-transport.md to use official kvcache-ai/Mooncake repo
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Docs] Update EFA benchmark results with tuned parameters (170 GB/s)
Update benchmark documentation with comprehensive parameter tuning results
from cross-machine testing on p6-b200.48xlarge instances. Key finding:
MC_SLICE_SIZE=262144 nearly doubles EFA throughput from ~70 to ~170 GB/s,
reaching 88% of RoCE RDMA performance.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* Fix clang-format violation in transfer_metadata.h
Remove extra space before comment on efa_addr field to satisfy
clang-format-20 style check.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Docs] Add EFA latency benchmark script, rename efa doc to underscore
- Add efa_latency_bench.py: automated benchmark script that measures
EFA throughput for tuned/default configs via SSH and plots
Latency vs Cache Size chart
- Add efa_latency_bench.png: benchmark results chart
- Rename efa-transport.md -> efa_transport.md to match naming
convention of other transport docs (ascend_transport.md, etc.)
- Update toctree reference in index.md
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---------
Co-authored-by: Ubuntu <ubuntu@ip-172-31-25-79.us-east-2.compute.internal>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: Ubuntu <ubuntu@ip-172-31-22-204.us-east-2.compute.internal>
This commit is contained in:
王鹤男
GitHub
parent
ab4af4f2da
commit
4136d2b73b
|
|
@ -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 <model_path> -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 <model_path> -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=<target_hostname>:<target_port> \
|
||||
--operation=write \
|
||||
--duration=10 \
|
||||
--threads=8 \
|
||||
--block_size=65536 \
|
||||
--batch_size=128 \
|
||||
--buffer_size=1073741824 \
|
||||
--report_unit=GB
|
||||
```
|
||||
|
||||
Replace `<target_hostname>:<target_port>` 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
|
||||
|
|
@ -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 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}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
$<TARGET_OBJECTS:rpc_communicator>
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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<TransferEngine>(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<TransferEngine>(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;
|
||||
}
|
||||
|
|
|
|||
Binary file not shown.
|
After Width: | Height: | Size: 63 KiB |
|
|
@ -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()
|
||||
|
|
@ -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<char *>(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") {
|
||||
|
|
|
|||
|
|
@ -117,6 +117,9 @@ class TransferMetadata {
|
|||
#endif
|
||||
std::vector<uint32_t> qp_num;
|
||||
std::string reply_msg; // on error
|
||||
#ifdef USE_EFA
|
||||
std::string efa_addr; // EFA endpoint address (hex encoded)
|
||||
#endif
|
||||
};
|
||||
|
||||
struct NotifyDesc {
|
||||
|
|
|
|||
|
|
@ -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 <gflags/gflags.h>
|
||||
#include <glog/logging.h>
|
||||
#include <rdma/fabric.h>
|
||||
#include <rdma/fi_domain.h>
|
||||
#include <rdma/fi_endpoint.h>
|
||||
#include <rdma/fi_cm.h>
|
||||
#include <rdma/fi_rma.h>
|
||||
#include <rdma/fi_errno.h>
|
||||
|
||||
#include <atomic>
|
||||
#include <condition_variable>
|
||||
#include <cstdint>
|
||||
#include <list>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <thread>
|
||||
#include <unordered_map>
|
||||
|
||||
#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<EfaEndPoint> 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<EfaEndPoint> getOrInsert(
|
||||
const std::string &peer_nic_path, std::shared_ptr<EfaEndPoint> new_ep);
|
||||
void add(const std::string &peer_nic_path,
|
||||
std::shared_ptr<EfaEndPoint> endpoint);
|
||||
void remove(const std::string &peer_nic_path);
|
||||
int disconnectAll();
|
||||
size_t size() const;
|
||||
|
||||
private:
|
||||
mutable RWSpinlock lock_;
|
||||
std::unordered_map<std::string, std::shared_ptr<EfaEndPoint>> 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<EfaEndPoint> 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<Transport::Slice *> &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<EfaEndpointStore> endpoint_store_;
|
||||
std::vector<std::shared_ptr<EfaCq>> cq_list_;
|
||||
|
||||
RWSpinlock mr_lock_;
|
||||
std::unordered_map<uint64_t, EfaMemoryRegionMeta> mr_map_;
|
||||
};
|
||||
|
||||
} // namespace mooncake
|
||||
|
||||
#endif // EFA_CONTEXT_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 <glog/logging.h>
|
||||
#include <rdma/fabric.h>
|
||||
#include <rdma/fi_domain.h>
|
||||
#include <rdma/fi_endpoint.h>
|
||||
#include <rdma/fi_cm.h>
|
||||
#include <rdma/fi_rma.h>
|
||||
|
||||
#include <atomic>
|
||||
#include <cstdint>
|
||||
#include <queue>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#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<Transport::Slice *> &slice_list,
|
||||
std::vector<Transport::Slice *> &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> 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<uint8_t> 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
|
||||
|
|
@ -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 <infiniband/verbs.h>
|
||||
|
||||
#include <atomic>
|
||||
#include <cstddef>
|
||||
#include <map>
|
||||
#include <memory>
|
||||
#include <mutex>
|
||||
#include <string>
|
||||
#include <thread>
|
||||
#include <unordered_map>
|
||||
#include <unordered_set>
|
||||
#include <vector>
|
||||
|
||||
#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<TransferMetadata> meta,
|
||||
std::shared_ptr<Topology> 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<BufferEntry> &buffer_list,
|
||||
const std::string &location) override;
|
||||
|
||||
int unregisterLocalMemoryBatch(
|
||||
const std::vector<void *> &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<TransferRequest> &entries) override;
|
||||
|
||||
Status submitTransferTask(
|
||||
const std::vector<TransferTask *> &task_list) override;
|
||||
|
||||
Status getTransferStatus(BatchID batch_id,
|
||||
std::vector<TransferStatus> &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<TransferMetadata> 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<std::shared_ptr<EfaContext>> context_list_;
|
||||
std::shared_ptr<Topology> local_topology_;
|
||||
|
||||
// CQ polling worker threads
|
||||
std::atomic<bool> worker_running_{false};
|
||||
std::vector<std::thread> worker_threads_;
|
||||
};
|
||||
|
||||
} // namespace mooncake
|
||||
|
||||
#endif // EFA_TRANSPORT_H_
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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 <cassert>
|
||||
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -56,3 +56,9 @@ if (USE_UBSHMEM)
|
|||
add_subdirectory(ascend_transport)
|
||||
target_sources(transport PUBLIC $<TARGET_OBJECTS:ascend_transport>)
|
||||
endif()
|
||||
|
||||
if (USE_EFA)
|
||||
add_subdirectory(efa_transport)
|
||||
target_sources(transport PUBLIC $<TARGET_OBJECTS:efa_transport>)
|
||||
target_link_libraries(transport PRIVATE fabric)
|
||||
endif()
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -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 <fcntl.h>
|
||||
#include <sys/epoll.h>
|
||||
|
||||
#include <atomic>
|
||||
#include <cassert>
|
||||
#include <fstream>
|
||||
#include <memory>
|
||||
#include <thread>
|
||||
#include <cstring>
|
||||
#include <iomanip>
|
||||
#include <sstream>
|
||||
|
||||
#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<EfaEndPoint> 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<EfaEndPoint> EfaEndpointStore::getOrInsert(
|
||||
const std::string &peer_nic_path, std::shared_ptr<EfaEndPoint> 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<EfaEndPoint> 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<EfaEndpointStore>();
|
||||
|
||||
// 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<EfaCq>();
|
||||
|
||||
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<EfaEndPoint> 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<EfaEndPoint>(*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<const uint8_t *>(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<Transport::Slice *> &slice_list) {
|
||||
// Route slices to appropriate endpoints for sending
|
||||
// Group slices by peer NIC path
|
||||
std::unordered_map<std::string, std::vector<Transport::Slice *>>
|
||||
slices_by_peer;
|
||||
std::vector<Transport::Slice *> 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<Transport::Slice *> 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<volatile int *, int> wr_depth_set;
|
||||
for (ssize_t i = 0; i < ret; i++) {
|
||||
EfaOpContext *op_ctx =
|
||||
reinterpret_cast<EfaOpContext *>(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<int>(ret));
|
||||
return static_cast<int>(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<volatile int *, int> wr_depth_set;
|
||||
|
||||
while ((ret = fi_cq_readerr(cq, &err_entry, 0)) > 0) {
|
||||
EfaOpContext *op_ctx =
|
||||
reinterpret_cast<EfaOpContext *>(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
|
||||
|
|
@ -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 <glog/logging.h>
|
||||
|
||||
#include <cassert>
|
||||
#include <cstddef>
|
||||
#include <cstring>
|
||||
#include <iomanip>
|
||||
#include <sstream>
|
||||
#include <thread>
|
||||
|
||||
#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<uint8_t> 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<Transport::Slice *> &slice_list,
|
||||
std::vector<Transport::Slice *> &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<int>(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
|
||||
|
|
@ -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 <glog/logging.h>
|
||||
#include <sys/mman.h>
|
||||
#include <sys/time.h>
|
||||
|
||||
#include <cassert>
|
||||
#include <chrono>
|
||||
#include <cstddef>
|
||||
#include <future>
|
||||
#include <set>
|
||||
#include <thread>
|
||||
|
||||
#include <dlfcn.h>
|
||||
|
||||
#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<TransferMetadata> meta,
|
||||
std::shared_ptr<Topology> 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<std::thread> threads;
|
||||
threads.reserve(num_threads);
|
||||
std::vector<int> thread_results(num_threads, 0);
|
||||
|
||||
for (size_t thread_i = 0; thread_i < num_threads; ++thread_i) {
|
||||
void *block_addr = static_cast<char *>(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<std::thread> reg_threads;
|
||||
reg_threads.reserve(context_list_.size());
|
||||
std::vector<int> 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<std::chrono::milliseconds>(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<MemoryLocationEntry> 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<std::thread> unreg_threads;
|
||||
unreg_threads.reserve(context_list_.size());
|
||||
std::vector<int> 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<SegmentDesc>();
|
||||
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<EfaTransport::BufferEntry> &buffer_list,
|
||||
const std::string &location) {
|
||||
std::vector<std::future<int>> 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<void *> &addr_list) {
|
||||
std::vector<std::future<int>> 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<TransferRequest> &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<TransferTask *> task_list;
|
||||
for (auto &task : batch_desc.task_list) task_list.push_back(&task);
|
||||
return submitTransferTask(task_list);
|
||||
}
|
||||
|
||||
Status EfaTransport::submitTransferTask(
|
||||
const std::vector<TransferTask *> &task_list) {
|
||||
std::unordered_map<std::shared_ptr<EfaContext>, std::vector<Slice *>>
|
||||
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<size_t>(device_id) < context_list_.size());
|
||||
auto &context = context_list_[device_id];
|
||||
assert(context.get());
|
||||
if (!context->active()) continue;
|
||||
assert(buffer_id >= 0 &&
|
||||
static_cast<size_t>(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<uintptr_t>(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<TransferStatus> &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<EfaContext> 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<std::string> efa_devices;
|
||||
std::vector<std::string> 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<EfaContext>(*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<int>(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
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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 <gflags/gflags.h>
|
||||
#include <glog/logging.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include <sys/time.h>
|
||||
|
||||
#include <cstdlib>
|
||||
#include <memory>
|
||||
|
||||
#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<TransferEngine> 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<TransferEngine>(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<TransferEngine>(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<TransferRequest> 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<TransferRequest> 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();
|
||||
}
|
||||
Loading…
Reference in New Issue