Compare commits

..

1 Commits

Author SHA1 Message Date
Teng Ma fa0641a436
[MISC] Add CODEOWNERS for efa_transport directory 2026-04-14 14:20:05 +08:00
150 changed files with 1625 additions and 12029 deletions

View File

@ -100,7 +100,7 @@ jobs:
sudo bash -x dependencies.sh -y
mkdir build
cd build
cmake -G Ninja .. -DUSE_HTTP=ON -DUSE_CXL=ON -DUSE_UB=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON -DENABLE_ASAN=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Debug
cmake -G Ninja .. -DUSE_HTTP=ON -DUSE_CXL=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON -DENABLE_ASAN=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Debug
shell: bash
- name: Build project
@ -134,12 +134,10 @@ jobs:
MASTER_PID=$!
sleep 3
cd mooncake-store/go
export LD_LIBRARY_PATH=$GITHUB_WORKSPACE/build/mooncake-common:$GITHUB_WORKSPACE/build/mooncake-store/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base:$GITHUB_WORKSPACE/build/mooncake-common/etcd
export LD_LIBRARY_PATH=$GITHUB_WORKSPACE/build/mooncake-asio:$GITHUB_WORKSPACE/build/mooncake-store/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base:$GITHUB_WORKSPACE/build/mooncake-common/etcd
export CGO_ENABLED=1
export CGO_CFLAGS="-I$GITHUB_WORKSPACE/mooncake-store/include -I$GITHUB_WORKSPACE/mooncake-transfer-engine/include"
export CGO_LDFLAGS="-L$GITHUB_WORKSPACE/build/mooncake-store/src -L$GITHUB_WORKSPACE/build/mooncake-store/src/cachelib_memory_allocator -L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src -L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base -L$GITHUB_WORKSPACE/build/mooncake-common -L$GITHUB_WORKSPACE/build/mooncake-common/etcd -lmooncake_store -lcachelib_memory_allocator -ltransfer_engine -lbase -lasio -letcd_wrapper -lstdc++ -lnuma -lglog -lgflags -libverbs -ljsoncpp -lzstd -lcurl -luring -lasan -lm -lgcov"
# Link cudart if CUDA is available (needed for D2H staging in mooncake_store)
if [ -d /usr/local/cuda/lib64 ]; then export CGO_LDFLAGS="$CGO_LDFLAGS -L/usr/local/cuda/lib64 -lcudart"; fi
export CGO_LDFLAGS="-L$GITHUB_WORKSPACE/build/mooncake-store/src -L$GITHUB_WORKSPACE/build/mooncake-store/src/cachelib_memory_allocator -L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src -L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base -L$GITHUB_WORKSPACE/build/mooncake-asio -L$GITHUB_WORKSPACE/build/mooncake-common/etcd -lmooncake_store -lcachelib_memory_allocator -ltransfer_engine -lbase -lasio -letcd_wrapper -lstdc++ -lnuma -lglog -lgflags -libverbs -ljsoncpp -lzstd -lcurl -luring -lasan -lm -lgcov"
ASAN_OPTIONS=detect_leaks=0:verify_asan_link_order=0 MC_METADATA_SERVER=http://127.0.0.1:8080/metadata go test -v ./tests/...
kill $MASTER_PID 2>/dev/null || true
shell: bash
@ -345,14 +343,6 @@ jobs:
- name: Run tests with ssd
run: |
# Reserve port 50052 (mooncake_client RPC port) so the kernel never
# auto-allocates it as ephemeral source port for other outbound
# connections in the test suite. Without this, a random Python test
# connection can pick src_port=50052, leave a TIME_WAIT on
# <eth0_ip>:50052 for 60s, and block mooncake_client's bind to
# 0.0.0.0:50052 even with SO_REUSEADDR (Linux only relaxes
# TIME_WAIT+bind conflict for same-IP or loopback).
sudo sysctl -w net.ipv4.ip_local_reserved_ports=50052
source test_env/bin/activate
MC_STORE_MEMCPY=false TEST_SSD_OFFLOAD_IN_EVICT=true ./scripts/run_tests.sh
rm -rf /tmp/mooncake_test_ssd
@ -518,7 +508,7 @@ jobs:
cd build
export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
cmake -G Ninja .. -DUSE_ETCD=OFF -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=OFF -DUSE_MNNVL=OFF -DUSE_UB=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs"
cmake -G Ninja .. -DUSE_ETCD=OFF -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=OFF -DUSE_MNNVL=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs"
cmake --build .
sudo cmake --install .
df -h
@ -528,7 +518,7 @@ jobs:
run: |
mkdir build
cd build
cmake -G Ninja .. -DUSE_ETCD=ON -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_STORE=ON -DWITH_P2P_STORE=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=ON -DUSE_MNNVL=OFF -DUSE_UB=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs"
cmake -G Ninja .. -DUSE_ETCD=ON -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_STORE=ON -DWITH_P2P_STORE=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=ON -DUSE_MNNVL=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs"
shell: bash
# TODO: lack USE_NVMEOF,USE_MNNVL
@ -746,18 +736,6 @@ jobs:
uses: ./.github/workflows/ci_cu13.yml
secrets: inherit
ascend-test:
needs: [build, check-paths]
if: needs.check-paths.outputs.should-run-downstream == 'true'
uses: ./.github/workflows/ci_ascend.yml
secrets: inherit
integration-test:
needs: [build, check-paths]
if: needs.check-paths.outputs.should-run-downstream == 'true'
uses: ./.github/workflows/integration-test.yml
secrets: inherit
ci-gate:
name: CI Gate
if: always()
@ -770,8 +748,6 @@ jobs:
- build-docker
- test-wheel-ubuntu
- build-wheel-cu13
- ascend-test
- integration-test
runs-on: ubuntu-latest
steps:
- name: Check required job results

View File

@ -15,8 +15,8 @@ jobs:
container:
image: localhost:5000/mooncake-hixl-ci:v5
options: --privileged --user 0:0 --device /dev/davinci0 --device /dev/davinci1 --device /dev/davinci2 --device /dev/davinci3
--device /dev/davinci4 --device /dev/davinci5 --device /dev/davinci6 --device /dev/davinci7
options: --privileged --user 0:0 --device /dev/davinci0 --device /dev/davinci1 --device /dev/davinci2 --device /dev/davinci3
--device /dev/davinci4 --device /dev/davinci5 --device /dev/davinci6 --device /dev/davinci7
--device /dev/davinci_manager --device /dev/devmm_svm --device /dev/hisi_hdc --ulimit nproc=65535:65535
env:
GITHUB_ACTIONS: "true"
@ -28,132 +28,31 @@ jobs:
- /etc/hccn.conf:/etc/hccn.conf
steps:
- name: Configure GitHub fetch defaults
shell: bash
run: |
git config --global protocol.version 2
git config --global http.version HTTP/1.1
git config --global http.lowSpeedLimit 1024
git config --global http.lowSpeedTime 30
- name: Checkout code
id: checkout_code
continue-on-error: true
uses: actions/checkout@v4
with:
ref: ${{ inputs.checkout_ref || github.sha }}
fetch-depth: 1
persist-credentials: false
- name: Retry checkout via GitHub mirror
if: steps.checkout_code.outcome == 'failure'
shell: bash
env:
ASCEND_GITHUB_MIRROR_URLS: ${{ vars.ASCEND_GITHUB_MIRROR_URLS }}
CHECKOUT_REF: ${{ inputs.checkout_ref || github.sha }}
run: |
set -euo pipefail
if [ -z "${ASCEND_GITHUB_MIRROR_URLS:-}" ]; then
echo "Checkout from GitHub failed and ASCEND_GITHUB_MIRROR_URLS is not set"
exit 1
fi
normalize_base() {
local base="$1"
base="${base#${base%%[![:space:]]*}}"
base="${base%${base##*[![:space:]]}}"
[ -n "$base" ] || return 1
[ "$base" != "https://github.com/" ] && base="${base%/}/"
printf '%s\n' "$base"
}
candidates=()
while IFS= read -r raw; do
base="$(normalize_base "$raw" || true)"
[ -n "$base" ] || continue
[ "$base" = "https://github.com/" ] && continue
candidates+=("$base")
done < <(printf '%s\n' "$ASCEND_GITHUB_MIRROR_URLS" | tr ',;' '\n')
if [ ${#candidates[@]} -eq 0 ]; then
echo "Checkout from GitHub failed and no valid mirror candidates were configured"
exit 1
fi
workdir="${GITHUB_WORKSPACE}"
git config --global --add safe.directory "$workdir"
for base in "${candidates[@]}"; do
mirror_url="${base}https://github.com/${GITHUB_REPOSITORY}.git"
echo "Retrying checkout with ${mirror_url}"
find "$workdir" -mindepth 1 -maxdepth 1 -exec rm -rf {} +
git init "$workdir"
git -C "$workdir" remote add origin "$mirror_url"
if git -C "$workdir" fetch --depth=1 origin "$CHECKOUT_REF" && \
git -C "$workdir" checkout --force --detach FETCH_HEAD; then
echo "Mirror checkout succeeded via ${base}"
exit 0
fi
echo "Mirror checkout failed via ${base}"
rm -rf "$workdir/.git"
done
echo "Direct GitHub checkout failed and all mirror retries failed"
exit 1
- name: Configure CMake
shell: bash
env:
ASCEND_GITHUB_MIRROR_URLS: ${{ vars.ASCEND_GITHUB_MIRROR_URLS }}
run: |
source /usr/local/Ascend/cann-9.0.0/set_env.sh
pwd
submodule_updated=false
if git submodule update --init --recursive; then
submodule_updated=true
elif [ -n "${ASCEND_GITHUB_MIRROR_URLS:-}" ]; then
normalize_base() {
local base="$1"
base="${base#${base%%[![:space:]]*}}"
base="${base%${base##*[![:space:]]}}"
[ -n "$base" ] || return 1
[ "$base" != "https://github.com/" ] && base="${base%/}/"
printf '%s\n' "$base"
}
while IFS= read -r raw; do
base="$(normalize_base "$raw" || true)"
[ -n "$base" ] || continue
[ "$base" = "https://github.com/" ] && continue
echo "Retrying submodule update with ${base}"
if git -c url."${base}https://github.com/".insteadOf=https://github.com/ \
submodule update --init --recursive; then
submodule_updated=true
break
fi
done < <(printf '%s\n' "$ASCEND_GITHUB_MIRROR_URLS" | tr ',;' '\n')
fi
if [ "$submodule_updated" != true ]; then
if [ ! -d "extern/pybind11" ] || [ -z "$(ls -A 'extern/pybind11' 2>/dev/null)" ]; then
echo "git submodule update failed (mirrors also exhausted), trying to cp pybind11..."
if [ -d "../pybind11" ]; then
cp -r ../pybind11 extern/
if ! git submodule update --init --recursive; then
if [ ! -d "extern/pybind11" ] || [ -z "$(ls -A 'extern/pybind11' 2>/dev/null)" ]; then
echo "git submodule update failed, try to cp pybind11..."
if [ -d "../pybind11" ]; then
cp -r ../pybind11 extern/
else
echo "Error: ../pybind11 does not exist. Cannot copy pybind11."
exit 1
fi
else
echo "Error: ../pybind11 does not exist. Cannot copy pybind11."
exit 1
echo "Detected that extern/pybind11 already exists, continuing execution...."
fi
else
echo "Detected that extern/pybind11 already exists, continuing execution...."
fi
fi
bash scripts/ascend/dependencies_ascend_installation.sh
echo "Configuring CMake..."
rm -rf build
@ -221,7 +120,7 @@ jobs:
# Check if master is running
if ! kill -0 $MASTER_PID 2>/dev/null; then
echo "Error: Mooncake Master failed to start"
cat /tmp/mooncake_master.log
cat /tmp/mooncake_master.log
exit 1
fi
@ -345,7 +244,7 @@ jobs:
echo ""
echo "All Hixl Mooncake Store tests completed successfully!"
- name: Test Summary
if: always()

1
.gitignore vendored
View File

@ -5,7 +5,6 @@ build_ofed4
old
local_test
go.sum
!mooncake-common/etcd/go.sum
*.so
bin
mod

View File

@ -23,16 +23,6 @@ repos:
- id: check-added-large-files
args: ['--maxkb=1024']
- repo: local
hooks:
- id: mooncake-code-format
name: Run Mooncake code format script
entry: ./scripts/code_format.sh
language: system
pass_filenames: false
always_run: true
require_serial: true
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.6.9
hooks:
@ -47,7 +37,7 @@ repos:
hooks:
- id: codespell
exclude: '^(extern/|FAST25-release/)'
args: ['--ignore-words-list=te,mooncake,KVCache,cann']
args: ['--ignore-words-list=te,mooncake,KVCache']
- repo: https://github.com/pre-commit/mirrors-clang-format
rev: v20.1.8

View File

@ -1,11 +1,10 @@
[default]
extend-ignore-words = ["CANN", "ASO", "fre", "wqs"]
extend-ignore-words = ["CANN", "ASO", "fre"]
[default.extend-words]
CANN = "CANN"
ASO = "ASO"
fre = "fre"
wqs = "wqs"
[files]
extend-exclude = [

View File

@ -45,21 +45,12 @@ option(STORE_USE_REDIS "build mooncake store with redis" OFF)
if (STORE_USE_REDIS)
add_compile_definitions(STORE_USE_REDIS)
endif()
option(STORE_USE_K8S_LEASE "build mooncake store with K8s Lease leader election" OFF)
if (STORE_USE_K8S_LEASE)
if (STORE_USE_ETCD)
message(FATAL_ERROR "STORE_USE_K8S_LEASE and STORE_USE_ETCD cannot be enabled together because both build Go c-shared HA backends.")
endif()
if (USE_ETCD AND NOT USE_ETCD_LEGACY)
message(FATAL_ERROR "STORE_USE_K8S_LEASE cannot be enabled with non-legacy USE_ETCD because both build Go c-shared libraries in the same process.")
endif()
add_compile_definitions(STORE_USE_K8S_LEASE)
endif()
option(STORE_USE_JEMALLOC "Use jemalloc in mooncake store master" OFF)
# Define ASIO macros before building targets that include ASIO headers.
# Define ASIO macros before adding mooncake-asio subdirectory
add_compile_definitions(ASIO_SEPARATE_COMPILATION ASIO_DYN_LINK)
add_subdirectory(mooncake-asio)
add_subdirectory(mooncake-common)
include_directories(mooncake-common/etcd)

View File

@ -41,7 +41,6 @@ Mooncake uses [pre-commit](https://pre-commit.com/) to enforce consistent format
| Type | Tool | Purpose |
|------|------|---------|
| Generic | trailing-whitespace / end-of-file-fixer | Basic hygiene |
| Project | `./scripts/code_format.sh` | Enforce Mooncake C/C++ formatting script before commit |
| Python | ruff / ruff-format | Lint + format (includes import sorting) |
| Spelling | codespell | Catch common typos (ignores domain-specific words) |
| C/C++ | clang-format | Apply style from the repository's `.clang-format` |
@ -54,8 +53,6 @@ pip install -r requirements-dev.txt
pre-commit install
```
After installation, every commit will run `./scripts/code_format.sh` automatically. If it rewrites files, re-stage the changes and commit again.
#### Usage
Run on all files (first run will install hook environments):
```bash

View File

@ -23,7 +23,7 @@ NC="\033[0m" # No Color
# Configuration
REPO_ROOT=`pwd`
GITHUB_PROXY=${GITHUB_PROXY:-"https://github.com"}
GOVER=1.25.9
GOVER=1.23.8
# Function to print section headers
print_section() {

View File

@ -88,11 +88,10 @@ store.setup(
| `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` | `/data/file_storage` | Absolute path to the SSD storage directory |
| `MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR` | `bucket_storage_backend` | Storage backend type (see below) |
| `MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES` | `1342177280` (1.25 GB) | Client-side staging buffer size |
| `MOONCAKE_OFFLOAD_SCANMETA_ITERATOR_KEYS_LIMIT` | `20000` | Max keys processed per iteration when scanning existing SSD metadata on startup |
| `MOONCAKE_OFFLOAD_TOTAL_SIZE_LIMIT_BYTES` | `2199023255552` (2 TB) | Maximum disk usage |
| `MOONCAKE_OFFLOAD_TOTAL_KEYS_LIMIT` | `10000000` | Maximum number of objects on disk |
| `MOONCAKE_OFFLOAD_HEARTBEAT_INTERVAL_SECONDS` | `10` | Interval for offload heartbeat to master (seconds) |
| `MOONCAKE_OFFLOAD_USE_URING` | `false` | Enable io_uring for async file I/O |
| `MOONCAKE_USE_URING` | `false` | Enable io_uring for async file I/O |
### Bucket backend settings
@ -102,8 +101,8 @@ Applies when `MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR=bucket_storage_backend
|---|---|---|
| `MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES` | `268435456` (256 MB) | Max size per bucket |
| `MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT` | `500` | Max keys per bucket |
| `MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE` | `0` | Eviction threshold in bytes. When set to `0`, the backend uses **90% of the physical disk capacity** as the quota — it does not mean unlimited. Set an explicit value to control disk usage precisely. |
| `MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY` | `none` | Eviction policy: `none` / `fifo` / `lru` |
| `MOONCAKE_BUCKET_MAX_TOTAL_SIZE` | `0` | Eviction threshold in bytes. When set to `0`, the backend uses **90% of the physical disk capacity** as the quota — it does not mean unlimited. Set an explicit value to control disk usage precisely. |
| `MOONCAKE_BUCKET_EVICTION_POLICY` | `none` | Eviction policy: `none` / `fifo` / `lru` |
---
@ -128,11 +127,6 @@ Best for: general-purpose use, large-scale deployments.
Stores each object in an individual file. Simple and easy to inspect, but generates many small files at scale.
| Environment Variable | Default | Description |
|---|---|---|
| `MOONCAKE_OFFLOAD_FSDIR` | `file_per_key_dir` | Subdirectory name under `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` where objects are stored |
| `MOONCAKE_OFFLOAD_ENABLE_EVICTION` | `true` | Enable disk eviction when the total size exceeds the quota |
Best for: debugging or small-scale deployments.
### `offset_allocator_storage_backend`
@ -149,7 +143,7 @@ Best for: high-concurrency scenarios with many small objects where restart durab
## Eviction (Bucket Backend Only)
When `MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE` is set, the backend automatically evicts buckets before writing new ones if total disk usage would exceed the limit.
When `MOONCAKE_BUCKET_MAX_TOTAL_SIZE` is set, the backend automatically evicts buckets before writing new ones if total disk usage would exceed the limit.
| Policy | Behavior |
|--------|----------|
@ -184,8 +178,8 @@ mooncake_master \
```bash
export MOONCAKE_OFFLOAD_FILE_STORAGE_PATH=/nvme/mooncake_offload
export MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR=bucket_storage_backend
export MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE=$((200 * 1024 * 1024 * 1024)) # 200 GB
export MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY=lru
export MOONCAKE_BUCKET_MAX_TOTAL_SIZE=$((200 * 1024 * 1024 * 1024)) # 200 GB
export MOONCAKE_BUCKET_EVICTION_POLICY=lru
mooncake_client \
--master_server_address="192.168.1.10:50051" \
@ -256,7 +250,7 @@ mooncake_client \
### "Failed to register buffer with UringFile" warning in logs
This warning appears when `MOONCAKE_OFFLOAD_USE_URING=true` and the io_uring fixed-buffer registration fails. The most common cause is that `MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES` exceeds the process's locked-memory limit (`RLIMIT_MEMLOCK`). io_uring requires the registered buffer to be pinned in physical memory, which counts against this limit.
This warning appears when `MOONCAKE_USE_URING=true` and the io_uring fixed-buffer registration fails. The most common cause is that `MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES` exceeds the process's locked-memory limit (`RLIMIT_MEMLOCK`). io_uring requires the registered buffer to be pinned in physical memory, which counts against this limit.
Check the current limit:

View File

@ -36,14 +36,13 @@ It is possible to configure a `Client` instance to act in only one of its two ro
* If `global_segment_size` is set to zero, the instance functions as a **pure client**, issuing requests but not contributing memory to the system.
* If `local_buffer_size` is set to zero, it acts as a **pure server**, providing memory for storage. In this case, request operations such as `Get` or `Put` are not permitted from this instance.
The `Client` can be used in three ways:
1. **Embedded mode**: Runs in the same process as the LLM inference program (e.g., a vLLM instance), by being imported as a shared library. Embedded clients issue requests directly, and when configured with `global_segment_size > 0` they also contribute memory resources to the cluster.
2. **Embedded mode with dummy-real clients**: Each LLM inference **rank** holds an embedded **dummy** client (which holds no resources). Each LLM inference **instance** has one resource-owning **real** client (for example, with TP=8 there can be 8 dummy clients and 1 real client). All dummy clients of the same inference instance forward requests to that one real client. The real client owns the global segment (optionally) and is responsible for RPC handling, memory management, and data transfer. Dummy and real clients communicate via RPC, and use shared memory/zero-copy mechanisms for data transfer, so that the data path remains efficient.
3. **Standalone store service**: A standalone store service (e.g., `python -m mooncake.mooncake_store_service`) wraps a client and provides the global memory/SSD resource pool. With this service, embedded clients can be configured with `global_segment_size = 0` so they contribute network/NIC resources only, while the standalone store service owns memory and storage management. This service can be deployed on the same server as the inference engine or on separate servers.
The `Client` can be used in two modes:
1. **Embedded mode**: Runs in the same process as the LLM inference program (e.g., a vLLM instance), by being imported as a shared library.
2. **Standalone mode**: Runs as an independent process. In this mode, the `Client` is separated into two parts: a **dummy** `Client` and a **real** `Client`: The **real** `Client` is a full-featured implementation that runs as a standalone process and directly communicates with other Mooncake Store components. It handles all RPC communications, memory management, and data transfer operations. The **real** `Client` is typically deployed on nodes that contribute memory to the distributed cache pool; The **dummy** `Client` is a lightweight wrapper that forwards all operations to a local **real** `Client` via RPC calls, which is designed for scenarios where the client needs to be embedded in the same process as the application (such as vLLM), but the actual Mooncake Store operations should be handled by a standalone process. The **dummy** `Client` and the **real** `Client` communicate via RPC calls and shared memory to make sure that Zero-copy transfers are still possible.
Mooncake store supports two deployment methods to accommodate different availability requirements:
1. **Default mode**: In this mode, the master service consists of a single master node, which simplifies deployment but introduces a single point of failure. If the master crashes or becomes unreachable, the system cannot continue to serve requests until it is restored.
2. **High availability mode**: This mode enhances fault tolerance by running the master service as a cluster of multiple master nodes coordinated through an etcd cluster. The master nodes use etcd to elect a leader, which is responsible for handling client requests.
2. **High availability mode (unstable)**: This mode enhances fault tolerance by running the master service as a cluster of multiple master nodes coordinated through an etcd cluster. The master nodes use etcd to elect a leader, which is responsible for handling client requests.
If the current leader fails or becomes partitioned from the network, the remaining master nodes automatically perform a new leader election, ensuring continuous availability.
In both modes, the leader monitors the health of all client nodes through periodic heartbeats. If a client crashes or becomes unreachable, the leader quickly detects the failure and takes appropriate action. When a client node recovers or reconnects, it can automatically rejoin the cluster without manual intervention.

View File

@ -154,7 +154,7 @@ A single pre-allocated file (`kv_cache.data`) is shared by all objects. Space wi
## Eviction (BucketStorageBackend)
When `MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE` is set, the backend evicts existing buckets to make room before writing a new one. Eviction is disabled by default (`BucketEvictionPolicy::NONE`).
When `MOONCAKE_BUCKET_MAX_TOTAL_SIZE` is set, the backend evicts existing buckets to make room before writing a new one. Eviction is disabled by default (`BucketEvictionPolicy::NONE`).
### Policies
@ -195,7 +195,7 @@ This ordering guarantees:
## io_uring File I/O
When `MOONCAKE_OFFLOAD_USE_URING=true`, the storage backends replace POSIX `pread`/`pwrite` calls with an io_uring-based implementation (`UringFile`). The design prioritizes eliminating inter-thread lock contention, which was the dominant latency source in the previous global-ring approach.
When `MOONCAKE_USE_URING=true`, the storage backends replace POSIX `pread`/`pwrite` calls with an io_uring-based implementation (`UringFile`). The design prioritizes eliminating inter-thread lock contention, which was the dominant latency source in the previous global-ring approach.
### Thread-local rings (`SharedUringRing`)

View File

@ -6,7 +6,7 @@ This document describes how to build and use Mooncake with AWS Elastic Fabric Ad
### 1. AWS EFA Driver and libfabric
EFA driver and libfabric should be pre-installed on AWS instances with EFA support (e.g., p6-b300.48xlarge, p6-b200.48xlarge, p5en.48xlarge, p5e.48xlarge, p5.48xlarge).
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
@ -22,30 +22,48 @@ If not installed, follow [AWS EFA documentation](https://docs.aws.amazon.com/AWS
### 2. Build Dependencies
Clone the repository and install all 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 \
libgtest-dev \
libmsgpack-dev \
libxxhash-dev \
libyaml-cpp-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
sudo ./dependencies.sh -y
git submodule update --init --recursive
```
This installs all system packages, git submodules (including pybind11 and yalantinglibs), and Go.
**Additional EFA-specific dependencies** (not covered by `dependencies.sh`):
```bash
# gflags is needed by transfer_engine_bench and EFA unit tests
sudo apt-get install -y libgflags-dev
```
> **Note:** The EFA driver and libfabric are **not** installed by `dependencies.sh`. They must be pre-installed on the instance (see section 1 above).
## Building Mooncake with EFA Support
### 1. Build with EFA Enabled
**GPU memory transfers (e.g., KV cache in vLLM):**
### 2. Build with EFA Enabled
```bash
mkdir build && cd build
@ -60,28 +78,13 @@ make -j$(nproc)
> **Note:** `-DUSE_CUDA=ON` is required when transferring GPU memory (e.g., KV cache in vLLM). Without it, the TCP transport (used as fallback when `mooncake_protocol` is set to `"tcp"`) cannot detect GPU memory and will fail with "Bad address" (EFAULT) errors.
**CPU memory transfers only (no GPU dependency):**
```bash
mkdir build && cd build
cmake .. \
-DUSE_EFA=ON \
-DUSE_CUDA=OFF \
-DCMAKE_BUILD_TYPE=RelWithDebInfo
make -j$(nproc)
```
> **Note:** With `-DUSE_CUDA=OFF`, the benchmark tool uses DRAM buffers allocated via `numa_alloc_onnode`. This is useful for measuring EFA transport throughput independently of GPU hardware.
### 2. Install Python Package
### 3. Install Python Package
```bash
# Copy built modules to wheel directory
cp mooncake-integration/engine.cpython-*.so ../mooncake-wheel/mooncake/
cp mooncake-integration/store.cpython-*.so ../mooncake-wheel/mooncake/
cp mooncake-common/libasio.so ../mooncake-wheel/mooncake/
cp mooncake-asio/libasio.so ../mooncake-wheel/mooncake/
# Install with pip
pip install -e ../mooncake-wheel --no-build-isolation
@ -102,6 +105,27 @@ print(f'Initialize result: {result}') # Should be 0
# 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):
@ -163,8 +187,6 @@ Use `transfer_engine_bench` to measure EFA transport throughput between two node
--report_unit=GB
```
> **Tip:** For CPU-to-CPU benchmarks, prepend `CUDA_VISIBLE_DEVICES=""` to prevent the CUDA runtime from being initialized. Without it, `nvidia-smi` may show GPU memory usage (due to CUDA context initialization) even though the benchmark only uses DRAM.
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
@ -174,219 +196,68 @@ Replace `<target_hostname>:<target_port>` with the target node's address shown i
| `--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 (per GPU when `--gpu_id=-1`) |
| `--buffer_size` | 1 GB | Total buffer size |
| `--duration` | 10 | Test duration in seconds |
| `--operation` | write | `read` or `write` |
| `--operation` | read | `read` or `write` |
| `--report_unit` | GB | `GB\|GiB\|Gb\|MB\|MiB\|Mb` |
| `--gpu_id` | 0 | GPU device ID; `-1` to use all GPUs (requires `-DUSE_CUDA=ON`) |
| Environment Variable | Default | Description |
|---------------------|---------|-------------|
| `MC_SLICE_SIZE` | 65536 | Slice size for RDMA transport. **Not used by EFA transport** (see note below). |
| `MC_EFA_STRIPING_THRESHOLD` | 2097152 | Transfers larger than this (bytes) are striped across all NICs |
> **Note on EFA slicing:** Unlike RDMA transport which splits every transfer into fixed `MC_SLICE_SIZE` chunks, EFA transport uses a different strategy: transfers ≤ `MC_EFA_STRIPING_THRESHOLD` (default 2MB) are sent as a **single `fi_write`/`fi_read`** whose size equals `block_size`; transfers larger than the threshold are striped across all NICs (one chunk per NIC). This means **`block_size` directly determines per-operation size** and is the key tuning parameter for EFA, while `MC_SLICE_SIZE` has no effect.
> **Note:** `buffer_size` must be >= `block_size * batch_size * threads`. The benchmark auto-adjusts if too small.
### Benchmark Results
#### p6-b200.48xlarge (B200, 8 EFA × 400 Gbps)
Tested on two p6-b200.48xlarge instances (8 EFA devices each, 8×400 Gbps) in the same AWS placement group.
Tested on two p6-b200.48xlarge instances in the same AWS placement group.
#### Optimized Results
**GPU-to-GPU** (build with `-DUSE_CUDA=ON`, `--gpu_id=-1` for all 8 GPUs):
With tuned parameters (`MC_SLICE_SIZE=262144`):
| Configuration | Write | Read |
|---------------|-------|------|
| block=1MB, threads=32, batch=64, buf=2GB/GPU | 285-296 GB/s | 312 GB/s |
| **block=1MB, threads=16, batch=128, buf=2GB/GPU** | **302 GB/s** | **313 GB/s** |
| 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 |
**CPU-to-CPU** (build with `-DUSE_CUDA=OFF`):
#### Parameter Tuning Results
| Configuration | Write | Read |
|---------------|-------|------|
| block=1MB, threads=32, batch=128, buf=4GB | **222 GB/s** (stable over 6 runs) | **226 GB/s** |
<details>
<summary>CPU Parameter Tuning History (p6-b200)</summary>
Earlier CPU-to-CPU tuning results (before EFA striping optimization, when `MC_SLICE_SIZE` was still used by EFA):
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 | 48 | 128 | 256KB | **160.34 GB/s** |
| 128KB | 64 | 128 | 256KB | 158.82 GB/s |
> **Note:** These results predate the EFA striping optimization. With the current code, `MC_SLICE_SIZE` no longer affects EFA performance. Use `--block_size=1048576` (1MB) instead, which achieves 222 GB/s.
</details>
#### p6-b300.48xlarge (B300, 16 EFA × 400 Gbps)
Tested on two p6-b300.48xlarge instances (Intel Xeon Platinum 8559C, 8× B300, 16 EFA devices) in the same AWS placement group.
**GPU-to-GPU** (build with `-DUSE_CUDA=ON`, `--gpu_id=-1` for all 8 GPUs, `--buffer_size=2147483648`):
| Configuration | Write | Read |
|---------------|-------|------|
| block=1MB, threads=16, batch=128 | 701 GB/s | **697 GB/s** |
| **block=1MB, threads=32, batch=64** | **752 GB/s** | 713 GB/s |
| block=1MB, threads=32, batch=32 | 751 GB/s | - |
| block=1MB, threads=64, batch=32 | 728 GB/s | - |
> **Peak: 752 GB/s write**, reaching ~94% of the 800 GB/s theoretical line rate (16×400 Gbps). GPUDirect RDMA bypasses DRAM entirely (HBM3e → PCIe switch → NIC), so performance is not bottlenecked by CPU memory bandwidth.
**CPU-to-CPU** (build with `-DUSE_CUDA=OFF`):
| Configuration | Write | Read |
|---------------|-------|------|
| **block=1MB, threads=32, batch=128, buf=4GB** | **230 GB/s** | 180 GB/s |
| block=16MB, threads=32, batch=8, buf=8GB (striping off) | 233 GB/s | - |
> CPU-to-CPU is bounded by DRAM bandwidth (~250 GB/s/socket on Xeon 8559C). Per-NIC sampling shows NUMA-0 NICs at 90 Gbps and NUMA-1 NICs at 53 Gbps, confirming DRAM controller saturation rather than NIC limit.
#### p5en.48xlarge (H200, 16 EFA × 200 Gbps)
Tested on two p5en.48xlarge instances (Intel Xeon 8488C, 8× H200 141GB, 16 EFA devices) in the same AWS placement group.
**GPU-to-GPU** (build with `-DUSE_CUDA=ON`, `--gpu_id=-1` for all 8 GPUs):
| Configuration | Write | Read |
|---------------|-------|------|
| block=1MB, threads=8, batch=128, buf=1GB/GPU | 236 GB/s | 271 GB/s |
| block=1MB, threads=16, batch=128, buf=2GB/GPU | 271 GB/s | **297-308 GB/s** |
| **block=1MB, threads=32, batch=64, buf=2GB/GPU** | **337-347 GB/s** | 274 GB/s |
> GPU HBM bandwidth (>3 TB/s) eliminates the memory bottleneck, allowing full EFA utilization. Write and read have different optimal thread counts: write peaks at 32 threads, read peaks at 16 threads.
> **Note:** EFA memory region registration (fi_mr_reg) for GPU memory segfaults at 4GB+ per GPU. Use `--buffer_size=2147483648` (2GB) as the maximum per-GPU buffer.
**CPU-to-CPU** (build with `-DUSE_CUDA=OFF`):
| Configuration | Write | Read |
|---------------|-------|------|
| Single instance (block=1MB, threads=32, batch=128, buf=4GB) | 179 GB/s | 185 GB/s |
| NUMA-split (block=1MB, 2 instances, 8 NICs each, threads=16, buf=2GB) | **192 GB/s** | **182 GB/s** |
> CPU-to-CPU throughput is bottlenecked by DRAM bandwidth (~155 GB/s per NUMA node, measured with STREAM Copy).
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 | Notes |
|-----------|-----------|-------|
| **EFA GPU-to-GPU (B300)** | **752 GB/s** | p6-b300.48xlarge, 16×400G, block=1MB, ~94% line rate |
| **EFA GPU-to-GPU (H200)** | **347 GB/s** | p5en.48xlarge, 16×200G, block=1MB |
| **EFA GPU-to-GPU (B200)** | **313 GB/s** | p6-b200.48xlarge, 8×400G, block=1MB |
| **EFA CPU-to-CPU (B300)** | **230 GB/s** | p6-b300.48xlarge, 16×400G, block=1MB, DRAM-limited |
| **EFA CPU-to-CPU (B200)** | **222 GB/s** | p6-b200.48xlarge, 8×400G, block=1MB, DRAM-limited |
| **EFA CPU-to-CPU (H200)** | **192 GB/s** | p5en.48xlarge, block=1MB, NUMA-split, DRAM-limited |
| EFA (default params) | 69.47 GB/s | Default block=64KB |
| TCP (iperf3 baseline) | 9.5 GB/s | Kernel TCP stack, 8 parallel streams |
| 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 vs RoCE RDMA**: On comparable 8×400 Gbps RoCE networks, Mooncake's RDMA transport achieves ~190 GB/s. Tuned EFA **exceeds** RoCE performance with GPU memory (313-347 GB/s) and on CPU-to-CPU (222 GB/s).
**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
- **Use `--block_size=1048576` (1MB)** — this is the most important tuning parameter for EFA. Each `block_size`-sized transfer becomes a single `fi_write`/`fi_read` call, so larger blocks amortize per-operation overhead. 1MB gives ~2× throughput over the 64KB default.
- `MC_SLICE_SIZE` has **no effect** on EFA transport (it only applies to RDMA transport). Use `block_size` instead.
- Increase `--threads` to 32-48 to saturate multiple EFA devices (2-4 threads per device is a good starting point)
- For **CPU-to-CPU**: use `--block_size=1048576` (1MB) with NUMA-split (separate instances per NUMA node) for best results
- For **GPU-to-GPU**: use `--block_size=1048576` (1MB), `--gpu_id=-1` (all GPUs), and `--buffer_size=2147483648` (2GB max per GPU). Write peaks at threads=32, read at threads=16
- Keep `--batch_size` such that `block_size * batch_size * threads <= buffer_size`
- Allocate buffers on both NUMA nodes for balanced NIC utilization (the bench tool does this by default for CPU mode)
- On 16-NIC instances (p5en), writes are NUMA-sensitive: 8 local-NUMA NICs reach 90 Gbps each, while 8 cross-NUMA NICs only reach ~20 Gbps without NUMA-split
### Eager endpoint warmup (first-request latency)
libfabric `FI_EP_RDM` endpoints resolve peer addresses lazily: `fi_av_insert()` and the metadata handshake fire on the first send to each `(local_ctx, peer_nic)` pair. On 16-NIC instances that gives `16 × N_peer_NICs` serial handshakes inside the first `submitTransfer`, which shows up as a single-digit-second first-batch stall (measured ~4 s on p6-B300 for a 100 × 0.5 MB batch; the first batch runs at <0.1 GB/s while the CQ drains, steady-state afterwards is unaffected).
Mooncake exposes an explicit eager-warmup API to eliminate the stall:
- C++: `EfaTransport::warmupSegment(const std::string& segment_name)`
- C: `int warmupEfaSegment(transfer_engine_t engine, const char *segment_name)`
- Rust: `TransferEngine::warmup_efa_segment(name: &str)`
Call it once per peer segment, right after `openSegment` (or after any metadata change that adds a new peer). Every `(local_ctx, peer_nic)` endpoint is connected concurrently via `std::async`; the critical path becomes `max(handshake RTT)` instead of `sum(handshake RTT)`. The call is idempotent — safe to re-run.
Measured on p6-B300 (16 local NICs × 16 peer NICs, dual-NUMA initiator, 100 × 0.5 MB batch):
| | first-batch latency | steady-state |
|---|---:|---:|
| No warmup | 4,043 ms | 141 GB/s |
| `warmup_efa_segment` (256 endpoints connected in 4.1 s) | **13.5 ms** (~300×) | 230 GB/s |
The warmup call itself takes roughly the same wall time as the stall it replaces — the win is that it's a one-time setup cost decoupled from the critical path of the first real transfer, not paid inside your latency budget.
## 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"}}'
```
## Usage with SGLang
SGLang's Mooncake integration currently hardcodes the `"rdma"` protocol. To use EFA transport, apply the provided patch and set environment variables.
### 1. Apply EFA Patch
SGLang's transfer engine initialization needs to be patched to read the protocol from an environment variable instead of using hardcoded `"rdma"`. Use the [patch script](https://github.com/whn09/kimi-k2-sglang):
```bash
bash patch_sglang_efa.sh
```
This is idempotent and safe to rerun.
### 2. Environment Variables
```bash
export MOONCAKE_PROTOCOL=efa
export FI_PROVIDER=efa
export FI_EFA_USE_DEVICE_RDMA=1
export GLOO_SOCKET_IFNAME=enp71s0 # adjust to your instance's primary interface
```
For multi-node expert parallelism (EP) deployments, also set:
```bash
export NVSHMEM_REMOTE_TRANSPORT=libfabric
export NVSHMEM_LIBFABRIC_PROVIDER=efa
```
> **Warning:** Do **not** set NVSHMEM variables on single-node deployments — doing so causes segmentation faults.
### 3. Docker Launch Example
```bash
docker run -d --name sglang \
--runtime=nvidia --gpus all --network host \
--privileged --shm-size=600g \
--device=/dev/infiniband \
-e MOONCAKE_PROTOCOL=efa \
-e FI_PROVIDER=efa \
-e FI_EFA_USE_DEVICE_RDMA=1 \
<image> bash start.sh
```
> **Note:** Ensure the Docker image's libfabric version matches the host's EFA driver. If not, mount the host's EFA libraries into the container (see [Troubleshooting](#libfabric-version-mismatch-in-docker)).
- **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
@ -419,11 +290,11 @@ AWS EFA exposes RDMA-like devices through the ibverbs interface, but does not su
### Thread Safety
The EFA transport requests `FI_THREAD_SAFE` from the libfabric provider and adds per-endpoint spinlocks to serialize `fi_write`/`fi_read` calls. This is necessary because:
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`/`fi_read` without serialization corrupts provider internals, causing completions to silently vanish
- 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.
@ -435,18 +306,14 @@ CQ completion queues are polled by dedicated worker threads (one per EFA device)
| 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 CPU-to-CPU (8×400G) | 222 GB/s (tuned) | ~190 GB/s |
| Throughput GPU-to-GPU (16×200G) | 347 GB/s (tuned) | N/A |
| Throughput GPU-to-GPU (8×400G) | 313 GB/s (tuned) | N/A |
| 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-b300.48xlarge (16 EFA devices × 400 Gbps = 6,400 Gbps, `rdmap*` naming)
- p6-b200.48xlarge (8 EFA devices × 400 Gbps = 3,200 Gbps, `rdmap*` naming)
- p5en.48xlarge (16 EFA devices × 200 Gbps = 3,200 Gbps, `rdmap*` naming)
- p5e.48xlarge (32 EFA devices × 100 Gbps = 3,200 Gbps, `rdmap*` naming)
- p5.48xlarge (32 EFA devices × 100 Gbps = 3,200 Gbps, `rdmap*` naming)
- 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.
@ -487,59 +354,3 @@ 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
### Building on AWS Deep Learning AMI
On AWS Deep Learning AMI (e.g., Ubuntu 24.04), the system Python and CUDA toolkit are bundled inside the `/opt/pytorch` virtual environment. You must activate it and set CUDA paths before building:
```bash
# Activate the PyTorch environment (provides Python 3.13 + CUDA toolkit)
source /opt/pytorch/bin/activate
# Set CUDA paths (nvcc, headers and libs are inside the pip-installed nvidia packages)
export CUDA_HOME=/opt/pytorch/lib/python3.13/site-packages/nvidia/cu13
export PATH=$CUDA_HOME/bin:$PATH
export CPLUS_INCLUDE_PATH=$CUDA_HOME/include:$CPLUS_INCLUDE_PATH
export LD_LIBRARY_PATH=$CUDA_HOME/lib:$LD_LIBRARY_PATH
export LIBRARY_PATH=$CUDA_HOME/lib:$LIBRARY_PATH
# Build with CUDA support
cd ~/Mooncake
mkdir -p build && cd build
cmake .. -DUSE_EFA=ON -DUSE_CUDA=ON -DCMAKE_BUILD_TYPE=RelWithDebInfo
make -j$(nproc)
```
Without activating the environment, you may encounter:
- `Could not find nvcc, please set CUDAToolkit_ROOT` — nvcc is not in PATH
- `fatal error: cuda.h: No such file or directory` — CUDA headers not in include path, set `CPLUS_INCLUDE_PATH`
- `cannot find -lcudart: No such file or directory` — CUDA libs not in library path, set `LIBRARY_PATH` and `LD_LIBRARY_PATH`
- `ModuleNotFoundError: No module named 'mooncake.engine'``.so` built against wrong Python version (e.g., 3.12 vs 3.13)
### libfabric version mismatch in Docker
```
fi_ep_bind (av) failed: Function not implemented
```
or:
```
undefined reference to `efadv_query_qp_wqs@EFA_1.4'
```
This happens when the Docker container's libfabric version is older than the host's EFA driver. Check with `fi_info --version` on both host and container.
Solution: Mount the host's EFA libraries into the container:
```bash
docker run --gpus all --device=/dev/infiniband --net=host --privileged \
-v /opt/amazon/efa:/opt/amazon/efa \
-v /lib/x86_64-linux-gnu/libefa.so.1:/lib/x86_64-linux-gnu/libefa.so.1 \
-v /lib/x86_64-linux-gnu/libefa.so:/lib/x86_64-linux-gnu/libefa.so \
-v /lib/x86_64-linux-gnu/libibverbs.so.1:/lib/x86_64-linux-gnu/libibverbs.so.1 \
-e LD_LIBRARY_PATH=/opt/amazon/efa/lib:$LD_LIBRARY_PATH \
-it <image>
```
Then rebuild Mooncake inside the container to link against the host's libfabric.

View File

@ -1,293 +0,0 @@
# Kunpeng UB Transport for Mooncake
This document describes how to build and use Mooncake with Kunpeng UB (Unified Bus) transport support using URMA (Unified Remote Memory Access).
## Overview
UB (Unified Bus) is a transport protocol at the same abstraction layer as RDMA, CXL, NVLink, and TCP, providing a flexible transport solution that can be selected at the application layer. Currently, UB protocol has two open-source implementations:
- **URMA (Unified Remote Memory Access)**: Provides a unified programming abstraction and core semantic layer for upper-layer applications. It offers unified APIs and semantic interfaces for remote shared memory access and operations, leveraging the low-latency, high-bandwidth characteristics of the UB protocol.
- URMA open-source repository: https://atomgit.com/openeuler/umdk
- **OBMM (Ownership Based Memory Management)**: A kernel memory management system for supernode environments, supporting cross-node physical memory sharing. It provides efficient remote memory access capabilities through a kernel module (obmm.ko) and a user-space library (libobmm.so).
- OBMM open-source repository: https://atomgit.com/openeuler/obmm
## Prerequisites
### 1. Hardware and Operating System
- **Hardware Platform**: Kunpeng 950 CPU with native UB interconnect architecture
- **OS Version**: openEuler 24.03 (LTS-SP3) [Download link](https://www.openeuler.openatom.cn/zh/download/#openEuler%2024.03%20LTS%20SP3)
### 2. URMA Dependencies
Install UMDK (URMA development package):
```bash
# Install via yum
yum install umdk-urma-devel
# Or build from source
git clone https://atomgit.com/openeuler/umdk.git
cd umdk
mkdir build && cd build
cmake ..
make -j$(nproc)
sudo make install
```
### 3. 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 \
libgtest-dev \
libmsgpack-dev \
libxxhash-dev \
libyaml-cpp-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 UB 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 UB Enabled
```bash
mkdir build && cd build
cmake .. \
-DUSE_UB=ON \
-DURMA_INCLUDE_DIR=/usr/include \
-DURMA_LIBRARY=/usr/lib64/liburma.so \
-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-integration/store.cpython-*.so ../mooncake-wheel/mooncake/
cp mooncake-common/libasio.so ../mooncake-wheel/mooncake/
# Install with pip
pip install -e ../mooncake-wheel --no-build-isolation
```
## Verification
### Check UB Transport Registration
```bash
# Check if UB transport is registered
./mooncake_server --list-transports
# Expected output: rdma, tcp, nvlink, ub
```
### Test UB Transport Initialization
```python
from mooncake.engine import TransferEngine
te = TransferEngine()
result = te.initialize('127.0.0.1', 'P2PHANDSHAKE', 'ub', '')
print(f'Initialize result: {result}') # Should be 0
# You should see logs like:
# URMA module init success
# found 1 devices.
# device_name : urma0 EID : 01:02:03:04:05:06:07:08:09:0a:0b:0c:0d:0e:0f:10
```
## Usage
### Single Node Benchmark Test
```bash
# Terminal 1: Target (receiver)
./transfer_engine_bench \
--mode=target \
--protocol=ub \
--device_name=urma0 \
--local_server_name=127.0.0.1 \
--metadata_server=P2PHANDSHAKE
# Terminal 2: Initiator (sender)
./transfer_engine_bench \
--mode=initiator \
--protocol=ub \
--device_name=urma0 \
--metadata_server=P2PHANDSHAKE \
--segment_size=8388608 \
--batch_size=1 \
--segment_id=127.0.0.1:$PORT
```
### Multi-device Benchmark Test
```bash
# Auto-discovery of multiple URMA devices
./transfer_engine_bench \
--protocol=ub \
--device_name=urma0,urma1,urma2,urma3
```
## Unit Tests
Run the UB transport unit tests:
```bash
./build/mooncake-transfer-engine/tests/ub_transport_test
```
The test suite includes:
| Test | Description |
|------|-------------|
| `MultiWrite` | Multiple write operations |
| `MultipleRead` | Multiple read operations with data integrity check |
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
```
## Technical Details
### UB Transport Architecture
```
┌─────────────────────────────────────────────────────┐
│ UbTransport │
├─────────────────────────────────────────────────────┤
│ UrmaContext (per device) │
│ ├── urma_device (URMA device handle) │
│ ├── urma_context (URMA context) │
│ ├── urma_jfce (URMA jetty factory create) │
│ ├── urma_jfc (URMA jetty factory send) │
│ └── urma_jfr (URMA jetty factory receive) │
├─────────────────────────────────────────────────────┤
│ UrmaEndpoint (per connection) │
│ ├── urma_jetty (URMA jetty for communication) │
│ ├── local_jetty (local jetty ID) │
│ └── remote_jetty (remote jetty ID) │
└─────────────────────────────────────────────────────┘
```
### Key Components
1. **UbTransport**: The main transport class that manages URMA resources and endpoints
2. **UrmaContext**: Represents a URMA device context, handling device initialization and resource management
3. **UrmaEndpoint**: Represents a connection to a remote peer, handling data transfer operations
4. **mock_urma_api.cpp**: Mock implementation of URMA API for testing without real URMA hardware
### Protocol Advantages
- **Optimized for Kunpeng**: URMA is specifically optimized for Kunpeng chip on-chip interconnect
- **RDMA-like Semantics**: Provides similar memory semantics to RDMA
- **High Performance**: Leverages UB's low-latency, high-bandwidth characteristics
- **Unified Abstraction**: Offers a unified programming model for remote memory access
## Troubleshooting
### No URMA devices found
```
UbTransport: No URMA devices found
```
Solution: Verify URMA is properly installed and devices are available:
```bash
# Check URMA installation
ls /usr/lib64/liburma.so
ls /usr/include/ub/umdk/urma/urma_api.h
# Check for URMA devices
urma_admin -l
```
### URMA initialization failed
```
URMA module init failed
```
Solution: Ensure the URMA kernel module is loaded and the device is properly configured:
```bash
# Load URMA module
sudo modprobe urma
# Check module status
sudo lsmod | grep urma
# Check device status
urma_admin -l
```
### Device port inactive
```
Device urma0 port not active
```
Solution: Ensure the UB port is properly configured and active:
```bash
# Check port status
urma_admin -p urma0
```
### Missing liburma.so
```
cannot find -lurma
```
Solution: Verify URMA library is installed and in the library path:
```bash
export LD_LIBRARY_PATH=/usr/lib64:$LD_LIBRARY_PATH
```
## Conclusion
Kunpeng UB Transport provides a high-performance, optimized transport solution for Mooncake on Kunpeng 950 CPU platforms. By leveraging the UB protocol's low-latency and high-bandwidth characteristics, it offers comparable performance to RDMA while being specifically tailored for Kunpeng chip architectures.
With proper configuration and tuning, UB Transport can significantly improve the performance of distributed AI workloads, particularly for scenarios involving large-scale parameter transfers and distributed training.

View File

@ -137,17 +137,13 @@ pip install mooncake-transfer-engine-non-cuda
make -j
```
5. If you want to compile MetaX (Muxi) MACA support (e.g. C500), install the MACA SDK so headers and libraries are available under `MACA_ROOT` (defaults to `MACA_HOME` env var if set, otherwise `/opt/maca`). SDK layouts vary; include both `lib` and `lib64` in runtime paths when needed:
5. If you want to compile MetaX (Muxi) MACA support (e.g. C500), install the MACA SDK so headers and libraries are available under `MACA_HOME` (default `/opt/maca` if unset). SDK layouts vary; CMake adds both `${MACA_HOME}/lib` and `${MACA_HOME}/lib64` to the link search path, so include both in the environment when linking or running binaries:
```bash
export MACA_HOME=/opt/maca
export LIBRARY_PATH=$LIBRARY_PATH:${MACA_HOME}/lib:${MACA_HOME}/lib64
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${MACA_HOME}/lib:${MACA_HOME}/lib64
```
Build with `-DUSE_MACA=ON`. Optional overrides:
- `-DMACA_ROOT=/path/to/maca`
- `-DMACA_INCLUDE_DIR=/path/to/maca/include`
- `-DMACA_LIB_DIR=/path/to/maca/lib64`
- `-DMACA_RUNTIME_LIBS="mcruntime;mxc-runtime64;rt"` (semicolon-separated CMake list)
Build with `-DUSE_MACA=ON`. To override which libraries are linked to `transfer_engine`, pass a CMake list at configure time, e.g. `-DMACA_RUNTIME_LIBS="mcruntime;mxc-runtime64;rt"` (semicolon-separated). This variable is read in `mooncake-transfer-engine/src/CMakeLists.txt` under `if(USE_MACA)`, not in `mooncake-common/common.cmake`.
6. Install yalantinglibs
```bash
@ -189,11 +185,7 @@ The following options can be used during `cmake ..` to specify whether to compil
- `-DUSE_CUDA=[ON|OFF]`: Enable GPU memory support (GPUDirect RDMA, NVMe-oF, and GPU-aware TCP transport). **Default: OFF.** Required when transferring GPU memory (e.g., KV cache in vLLM disaggregated serving), even when using TCP protocol.
- `-DUSE_MNNVL=[ON|OFF]`: Enable Multi-Node NVLink transport support, default is OFF. **Note:** `-DUSE_CUDA` is required when `-DUSE_MNNVL` is on (not used when building with `-DUSE_MUSA=ON`, `-DUSE_HIP=ON`, or `-DUSE_MACA=ON`).
- `-DUSE_MUSA=[ON|OFF]`: Enable Moore Threads GPU support via MUSA
- `-DUSE_MACA=[ON|OFF]`: Enable MetaX (Muxi) GPU support via MACA.
- `-DMACA_ROOT=/path/to/maca`: Override the MACA SDK root (`MACA_HOME` env var is also honored; default `/opt/maca`).
- `-DMACA_INCLUDE_DIR=/path/to/include`: Override MACA include directory when `-DUSE_MACA=ON`.
- `-DMACA_LIB_DIR=/path/to/lib64`: Override MACA library directory when `-DUSE_MACA=ON`.
- `-DMACA_RUNTIME_LIBS="mcruntime;mxc-runtime64;rt"`: Override MACA runtime libraries linked by `transfer_engine`.
- `-DUSE_MACA=[ON|OFF]`: Enable MetaX (Muxi) GPU support via MACA. Use `MACA_HOME` for the SDK root (default `/opt/maca`). Optional: `-DMACA_RUNTIME_LIBS` (CMake list) overrides default `mcruntime;mxc-runtime64;rt` for `transfer_engine` in `mooncake-transfer-engine/src/CMakeLists.txt`.
- `-DUSE_HIP=[ON|OFF]`: Enable AMD GPU support via HIP/ROCm
- `-DUSE_MLU=[ON|OFF]`: Enable Cambricon MLU memory support via Neuware. **Default: OFF.** Supports MLU memory detection, topology discovery, and RDMA registration for Transfer Engine.
- `-DNEUWARE_ROOT=/path/to/neuware`: Override the default Neuware SDK root used when `-DUSE_MLU=ON`. If unset, Mooncake uses `NEUWARE_HOME` or `/usr/local/neuware`.

View File

@ -141,7 +141,7 @@ mooncake_master \
# SSD Offload Buffer: 20GB
MOONCAKE_OFFLOAD_FILE_STORAGE_PATH="/mnt/data/file_storage" \
MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES=21474836480 \
MOONCAKE_OFFLOAD_USE_URING=1 \
MOONCAKE_USE_URING=1 \
mooncake_client \
--host=127.0.0.1 \
--global_segment_size=80GB \

View File

@ -132,17 +132,13 @@
make -j
```
5. 若需编译沐曦 MetaX MACA 支持(如 C500请安装 MACA SDK使头文件与库位于 `MACA_ROOT`(优先取 `MACA_HOME` 环境变量,未设置时默认 `/opt/maca`)。不同安装包可能把库放在 `lib``lib64`,建议在环境变量中同时加入两者,避免链接或运行时找不到共享库:
5. 若需编译沐曦 MetaX MACA 支持(如 C500请安装 MACA SDK使头文件与库位于 `MACA_HOME`未设置时默认 `/opt/maca`)。不同安装包可能把库放在 `lib``lib64`CMake 会同时搜索 `${MACA_HOME}/lib``${MACA_HOME}/lib64`,建议在环境变量中同时加入两者,避免链接或运行时找不到共享库:
```bash
export MACA_HOME=/opt/maca
export LIBRARY_PATH=$LIBRARY_PATH:${MACA_HOME}/lib:${MACA_HOME}/lib64
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${MACA_HOME}/lib:${MACA_HOME}/lib64
```
使用 `-DUSE_MACA=ON` 配置构建。可选覆盖项:
- `-DMACA_ROOT=/path/to/maca`
- `-DMACA_INCLUDE_DIR=/path/to/maca/include`
- `-DMACA_LIB_DIR=/path/to/maca/lib64`
- `-DMACA_RUNTIME_LIBS="mcruntime;mxc-runtime64;rt"`(分号分隔的 CMake 列表)
使用 `-DUSE_MACA=ON` 配置构建。如需覆盖默认链接的运行库,可在配置阶段传入 CMake 列表,例如 `-DMACA_RUNTIME_LIBS="mcruntime;mxc-runtime64;rt"`(分号分隔)。该变量在 **`mooncake-transfer-engine/src/CMakeLists.txt`** 的 `if(USE_MACA)` 中用于 `target_link_libraries`**不在** `mooncake-common/common.cmake` 中。
6. 安装 yalantinglibs
```bash
@ -171,11 +167,7 @@
在执行 `cmake ..` 期间可以使用下列选项指定是否编译 Mooncake 的某些组件。
- `-DUSE_CUDA=[ON|OFF]`: 启用 GPU Direct RDMA 及 NVMe-of 支持
- `-DUSE_MUSA=[ON|OFF]`: 通过 MUSA 启用对摩尔线程 GPU 的支持
- `-DUSE_MACA=[ON|OFF]`: 通过 MACA 启用对沐曦 MetaX GPU 的支持。
- `-DMACA_ROOT=/path/to/maca`: 覆盖 MACA SDK 根路径(也支持 `MACA_HOME` 环境变量,默认 `/opt/maca`)。
- `-DMACA_INCLUDE_DIR=/path/to/include`: 在 `-DUSE_MACA=ON` 时覆盖 MACA 头文件目录。
- `-DMACA_LIB_DIR=/path/to/lib64`: 在 `-DUSE_MACA=ON` 时覆盖 MACA 库目录。
- `-DMACA_RUNTIME_LIBS="mcruntime;mxc-runtime64;rt"`: 覆盖 `transfer_engine` 链接的 MACA 运行时库列表。
- `-DUSE_MACA=[ON|OFF]`: 通过 MACA 启用对沐曦 MetaX GPU 的支持(`MACA_HOME` 指向 SDK 根目录,默认 `/opt/maca`;可选 `-DMACA_RUNTIME_LIBS` 覆盖 `transfer_engine` 默认链接库,见 `mooncake-transfer-engine/src/CMakeLists.txt`
- `-DUSE_MLU=[ON|OFF]`: 通过 Neuware 启用寒武纪 MLU 显存支持。默认 OFF支持 MLU 显存探测、拓扑发现及 Transfer Engine 的 RDMA 注册。
- `-DNEUWARE_ROOT=/path/to/neuware`: 在 `-DUSE_MLU=ON` 时覆盖默认 Neuware SDK 根路径;未设置时使用 `NEUWARE_HOME``/usr/local/neuware`
- `-DMLU_INCLUDE_DIR=/path/to/include` / `-DMLU_LIB_DIR=/path/to/lib64`: 在 `-DUSE_MLU=ON` 时覆盖 Neuware 头文件与库目录。

View File

@ -0,0 +1,50 @@
# Build asio as a shared library to avoid ODR violations
# when multiple shared libraries use asio
# Try to find ASIO using find_package first
find_package(asio QUIET)
if(asio_FOUND)
message(STATUS "Found ASIO via find_package")
set(ASIO_INCLUDE_DIR ${asio_INCLUDE_DIR})
else()
# Fallback to find_path if find_package fails
find_path(ASIO_INCLUDE_DIR
NAMES asio.hpp
PATHS
/usr/local/include
/usr/include
${CMAKE_INSTALL_PREFIX}/include
DOC "Path to ASIO headers"
)
if(NOT ASIO_INCLUDE_DIR)
message(FATAL_ERROR "ASIO not found. Please install ASIO or set ASIO_INCLUDE_DIR manually.")
endif()
message(STATUS "Found ASIO at: ${ASIO_INCLUDE_DIR}")
endif()
add_library(asio_shared SHARED asio_impl.cpp)
target_compile_definitions(asio_shared
PUBLIC
ASIO_SEPARATE_COMPILATION
ASIO_DYN_LINK
)
target_include_directories(asio_shared
PUBLIC
${ASIO_INCLUDE_DIR}
)
set_target_properties(asio_shared PROPERTIES
POSITION_INDEPENDENT_CODE ON
INSTALL_RPATH "$ORIGIN"
BUILD_WITH_INSTALL_RPATH TRUE
OUTPUT_NAME "asio"
)
target_link_libraries(asio_shared PUBLIC pthread)
install(TARGETS asio_shared DESTINATION lib)

View File

@ -2,10 +2,6 @@ if ((USE_ETCD AND NOT USE_ETCD_LEGACY) OR STORE_USE_ETCD)
add_subdirectory(etcd)
endif()
if (STORE_USE_K8S_LEASE)
add_subdirectory(k8s-lease)
endif()
include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include)
add_subdirectory(src)

View File

@ -1,20 +0,0 @@
include(FetchContent)
# UMDK
FetchContent_Declare(
urma
GIT_REPOSITORY https://atomgit.com/openeuler/umdk.git
GIT_TAG v25.12.0
)
FetchContent_MakeAvailable(urma)
#
message(STATUS "URMA source dir: ${urma_SOURCE_DIR}")
message(STATUS "URMA binary dir: ${urma_BINARY_DIR}")
# UMDK include
set(urma_INCLUDE_DIR ${urma_SOURCE_DIR}/src/urma/lib/urma/core/include)
#
message(STATUS "urma_INCLUDE_DIR: ${urma_INCLUDE_DIR}")

View File

@ -79,9 +79,7 @@ option(USE_UB "option for using UB protocol transport" OFF)
if (USE_UB)
add_compile_definitions(USE_UB)
message(STATUS "ub transport is enabled")
include(${CMAKE_CURRENT_LIST_DIR}/FindUrma.cmake)
endif()
if (USE_EFA)
# Find libfabric headers and library; default to AWS EFA installer path
find_path(LIBFABRIC_INCLUDE_DIR rdma/fabric.h
@ -175,26 +173,6 @@ if (NOT DEFINED MLU_LIB_DIR OR MLU_LIB_DIR STREQUAL "")
set(MLU_LIB_DIR "${NEUWARE_ROOT}/lib64")
endif()
if (NOT DEFINED MACA_ROOT OR MACA_ROOT STREQUAL "")
if (DEFINED ENV{MACA_HOME} AND NOT "$ENV{MACA_HOME}" STREQUAL "")
set(MACA_ROOT "$ENV{MACA_HOME}" CACHE PATH "Path to MACA SDK" FORCE)
else()
set(MACA_ROOT "/opt/maca" CACHE PATH "Path to MACA SDK" FORCE)
endif()
endif()
if (NOT DEFINED MACA_INCLUDE_DIR OR MACA_INCLUDE_DIR STREQUAL "")
set(MACA_INCLUDE_DIR "${MACA_ROOT}/include")
endif()
if (NOT DEFINED MACA_LIB_DIR OR MACA_LIB_DIR STREQUAL "")
if (EXISTS "${MACA_ROOT}/lib64")
set(MACA_LIB_DIR "${MACA_ROOT}/lib64")
else()
set(MACA_LIB_DIR "${MACA_ROOT}/lib")
endif()
endif()
if (USE_MLU)
add_compile_definitions(USE_MLU)
message(STATUS "MLU support is enabled")
@ -205,12 +183,20 @@ if (USE_MLU)
endif()
if (USE_MACA)
# MACA toolchain is CUDA-compatible in first-stage porting.
# Reuse CUDA code paths to get a runnable baseline quickly.
add_compile_definitions(USE_MACA)
message(STATUS "MACA support is enabled")
include_directories(${MACA_INCLUDE_DIR})
if (EXISTS "${MACA_LIB_DIR}")
link_directories(${MACA_LIB_DIR})
if(DEFINED ENV{MACA_HOME})
set(MACA_HOME $ENV{MACA_HOME})
else()
set(MACA_HOME /opt/maca)
endif()
include_directories(${MACA_HOME}/include)
link_directories(
${MACA_HOME}/lib
${MACA_HOME}/lib64
)
endif()
if (USE_MUSA)

View File

@ -3,7 +3,7 @@ add_custom_command(
COMMAND bash -c "go mod tidy" && bash -c "go build -buildmode=c-shared -o ${CMAKE_CURRENT_BINARY_DIR}/libetcd_wrapper.so etcd_wrapper.go" && cp ${CMAKE_CURRENT_BINARY_DIR}/libetcd_wrapper.h ${CMAKE_CURRENT_SOURCE_DIR}
WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
COMMENT "Building Go shared library"
DEPENDS etcd_wrapper.go go.mod go.sum build.sh
DEPENDS etcd_wrapper.go
)
set(ETCD_WRAPPER_INCLUDE ${CMAKE_CURRENT_BINARY_DIR}/libetcd_wrapper.h)
@ -17,4 +17,4 @@ add_custom_target(
install(
FILES ${ETCD_WRAPPER_LIB}
DESTINATION lib
)
)

View File

@ -1,8 +1,6 @@
module github.com/kvcache-ai/Mooncake/mooncake-common/etcd
go 1.25.0
toolchain go1.25.9
go 1.24.0
require (
go.etcd.io/etcd/api/v3 v3.5.21

View File

@ -1,108 +0,0 @@
github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
github.com/coreos/go-semver v0.3.0 h1:wkHLiw0WNATZnSG7epLsujiMCgPAc9xhjJ4tgnAxmfM=
github.com/coreos/go-semver v0.3.0/go.mod h1:nnelYz7RCh+5ahJtPPxZlU+153eP4D4r3EedlOD2RNk=
github.com/coreos/go-systemd/v22 v22.3.2 h1:D9/bQk5vlXQFZ6Kwuu6zaiXJ9oTPe68++AzAJc1DzSI=
github.com/coreos/go-systemd/v22 v22.3.2/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI=
github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8=
github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
github.com/pkg/errors v0.8.1 h1:iURUrRGxPUNPdy5/HRSm+Yj6okJ6UtLINN0Q9M4+h3I=
github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
go.etcd.io/etcd/api/v3 v3.5.21 h1:A6O2/JDb3tvHhiIz3xf9nJ7REHvtEFJJ3veW3FbCnS8=
go.etcd.io/etcd/api/v3 v3.5.21/go.mod h1:c3aH5wcvXv/9dqIw2Y810LDXJfhSYdHQ0vxmP3CCHVY=
go.etcd.io/etcd/client/pkg/v3 v3.5.21 h1:lPBu71Y7osQmzlflM9OfeIV2JlmpBjqBNlLtcoBqUTc=
go.etcd.io/etcd/client/pkg/v3 v3.5.21/go.mod h1:BgqT/IXPjK9NkeSDjbzwsHySX3yIle2+ndz28nVsjUs=
go.etcd.io/etcd/client/v3 v3.5.21 h1:T6b1Ow6fNjOLOtM0xSoKNQt1ASPCLWrF9XMHcH9pEyY=
go.etcd.io/etcd/client/v3 v3.5.21/go.mod h1:mFYy67IOqmbRf/kRUvsHixzo3iG+1OF2W2+jVIQRAnU=
go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64=
go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y=
go.opentelemetry.io/otel v1.39.0 h1:8yPrr/S0ND9QEfTfdP9V+SiwT4E0G7Y5MO7p85nis48=
go.opentelemetry.io/otel v1.39.0/go.mod h1:kLlFTywNWrFyEdH0oj2xK0bFYZtHRYUdv1NklR/tgc8=
go.opentelemetry.io/otel/metric v1.39.0 h1:d1UzonvEZriVfpNKEVmHXbdf909uGTOQjA0HF0Ls5Q0=
go.opentelemetry.io/otel/metric v1.39.0/go.mod h1:jrZSWL33sD7bBxg1xjrqyDjnuzTUB0x1nBERXd7Ftcs=
go.opentelemetry.io/otel/sdk v1.39.0 h1:nMLYcjVsvdui1B/4FRkwjzoRVsMK8uL/cj0OyhKzt18=
go.opentelemetry.io/otel/sdk v1.39.0/go.mod h1:vDojkC4/jsTJsE+kh+LXYQlbL8CgrEcwmt1ENZszdJE=
go.opentelemetry.io/otel/sdk/metric v1.39.0 h1:cXMVVFVgsIf2YL6QkRF4Urbr/aMInf+2WKg+sEJTtB8=
go.opentelemetry.io/otel/sdk/metric v1.39.0/go.mod h1:xq9HEVH7qeX69/JnwEfp6fVq5wosJsY1mt4lLfYdVew=
go.opentelemetry.io/otel/trace v1.39.0 h1:2d2vfpEDmCJ5zVYz7ijaJdOF59xLomrvj7bjt6/qCJI=
go.opentelemetry.io/otel/trace v1.39.0/go.mod h1:88w4/PnZSazkGzz/w84VHpQafiU4EtqqlVdxWy+rNOA=
go.uber.org/atomic v1.7.0 h1:ADUqmZGgLDDfbSL9ZmPxKTybcoEYHgpYfELNoN+7hsw=
go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
go.uber.org/multierr v1.6.0 h1:y6IPFStTAIT5Ytl7/XYmHvzXQ7S3g/IeZW9hyZ5thw4=
go.uber.org/multierr v1.6.0/go.mod h1:cdWPpRnG4AhwMwsgIHip0KRBQjJy5kYEpYjJxpXp9iU=
go.uber.org/zap v1.17.0 h1:MTjgFu6ZLKvY6Pvaqk97GlxNBuMpV4Hy/3P6tRGlI2U=
go.uber.org/zap v1.17.0/go.mod h1:MXVU+bhUf/A7Xi2HNOnopQOrmycQ5Ih87HtOu4q5SSo=
golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
golang.org/x/net v0.48.0 h1:zyQRTTrjc33Lhh0fBgT/H3oZq9WuvRR5gPC70xpDiQU=
golang.org/x/net v0.48.0/go.mod h1:+ndRgGjkh8FGtu1w1FGbEC31if4VrNVMuKTgcAAnQRY=
golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.39.0 h1:CvCKL8MeisomCi6qNZ+wbb0DN9E5AATixKsvNtMoMFk=
golang.org/x/sys v0.39.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
golang.org/x/text v0.32.0 h1:ZD01bjUt1FQ9WJ0ClOL5vxgxOI/sVCNgX1YtKwcY0mU=
golang.org/x/text v0.32.0/go.mod h1:o/rUWzghvpD5TXrTIBuJU77MTaN0ljMWE47kxGJQ7jY=
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE=
golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk=
gonum.org/v1/gonum v0.16.0/go.mod h1:fef3am4MQ93R2HHpKnLk4/Tbh/s0+wqD5nfa6Pnwy4E=
google.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217 h1:fCvbg86sFXwdrl5LgVcTEvNC+2txB5mgROGmRL5mrls=
google.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217/go.mod h1:+rXWjjaukWZun3mLfjmVnQi18E1AsFbDN9QdJ5YXLto=
google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217 h1:gRkg/vSppuSQoDjxyiGfN4Upv/h/DQmIR10ZU8dh4Ww=
google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217/go.mod h1:7i2o+ce6H/6BluujYR+kqX3GKH+dChPTQU19wjRPiGk=
google.golang.org/grpc v1.79.3 h1:sybAEdRIEtvcD68Gx7dmnwjZKlyfuc61Dyo9pGXXkKE=
google.golang.org/grpc v1.79.3/go.mod h1:KmT0Kjez+0dde/v2j9vzwoAScgEPx/Bw1CYChhHLrHQ=
google.golang.org/protobuf v1.36.10 h1:AYd7cD/uASjIL6Q9LiTjz8JLcrh/88q5UObnmY3aOOE=
google.golang.org/protobuf v1.36.10/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

View File

@ -1,20 +0,0 @@
add_custom_command(
OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/libk8s_lease_wrapper.so
COMMAND bash -c "go mod tidy" && bash -c "go build -buildmode=c-shared -o ${CMAKE_CURRENT_BINARY_DIR}/libk8s_lease_wrapper.so k8s_lease_wrapper.go" && cp ${CMAKE_CURRENT_BINARY_DIR}/libk8s_lease_wrapper.h ${CMAKE_CURRENT_SOURCE_DIR}
WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
COMMENT "Building K8s Lease Go shared library"
DEPENDS k8s_lease_wrapper.go
)
set(K8S_LEASE_WRAPPER_INCLUDE ${CMAKE_CURRENT_BINARY_DIR}/libk8s_lease_wrapper.h)
set(K8S_LEASE_WRAPPER_LIB ${CMAKE_CURRENT_BINARY_DIR}/libk8s_lease_wrapper.so)
add_custom_target(
build_k8s_lease_wrapper
DEPENDS ${K8S_LEASE_WRAPPER_LIB}
)
install(
FILES ${K8S_LEASE_WRAPPER_LIB}
DESTINATION lib
)

View File

@ -1,61 +0,0 @@
// envtest-server starts a real kube-apiserver + etcd via envtest, writes the
// KUBECONFIG path to stdout, and blocks until SIGTERM or SIGINT. This lets
// C++ tests launch it as a subprocess and talk to a real K8s API without a
// full cluster.
package main
import (
"fmt"
"os"
"os/signal"
"path/filepath"
"syscall"
"k8s.io/client-go/tools/clientcmd"
clientcmdapi "k8s.io/client-go/tools/clientcmd/api"
"sigs.k8s.io/controller-runtime/pkg/envtest"
)
func main() {
env := &envtest.Environment{}
cfg, err := env.Start()
if err != nil {
fmt.Fprintf(os.Stderr, "envtest start failed: %v\n", err)
os.Exit(1)
}
// Write a KUBECONFIG file that points at the envtest kube-apiserver.
kubeconfigPath := filepath.Join(os.TempDir(), fmt.Sprintf("envtest-kubeconfig-%d", os.Getpid()))
kubeconfig := clientcmdapi.NewConfig()
kubeconfig.Clusters["envtest"] = &clientcmdapi.Cluster{
Server: cfg.Host,
CertificateAuthorityData: cfg.CAData,
}
kubeconfig.AuthInfos["envtest"] = &clientcmdapi.AuthInfo{
ClientCertificateData: cfg.CertData,
ClientKeyData: cfg.KeyData,
}
kubeconfig.Contexts["envtest"] = &clientcmdapi.Context{
Cluster: "envtest",
AuthInfo: "envtest",
}
kubeconfig.CurrentContext = "envtest"
if err := clientcmd.WriteToFile(*kubeconfig, kubeconfigPath); err != nil {
fmt.Fprintf(os.Stderr, "failed to write kubeconfig: %v\n", err)
env.Stop()
os.Exit(1)
}
// Print the kubeconfig path — the parent process reads this from stdout.
fmt.Println(kubeconfigPath)
// Block until SIGTERM or SIGINT.
sigCh := make(chan os.Signal, 1)
signal.Notify(sigCh, syscall.SIGTERM, syscall.SIGINT)
<-sigCh
os.Remove(kubeconfigPath)
env.Stop()
}

View File

@ -1,60 +0,0 @@
module github.com/kvcache-ai/Mooncake/mooncake-common/k8s-lease
go 1.24.0
require (
k8s.io/api v0.34.3
k8s.io/apimachinery v0.34.3
k8s.io/client-go v0.34.3
k8s.io/utils v0.0.0-20251002143259-bc988d571ff4
sigs.k8s.io/controller-runtime v0.22.5
)
require (
github.com/beorn7/perks v1.0.1 // indirect
github.com/cespare/xxhash/v2 v2.3.0 // indirect
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
github.com/emicklei/go-restful/v3 v3.12.2 // indirect
github.com/evanphx/json-patch/v5 v5.9.11 // indirect
github.com/fxamacker/cbor/v2 v2.9.0 // indirect
github.com/go-logr/logr v1.4.3 // indirect
github.com/go-openapi/jsonpointer v0.21.0 // indirect
github.com/go-openapi/jsonreference v0.20.2 // indirect
github.com/go-openapi/swag v0.23.0 // indirect
github.com/gogo/protobuf v1.3.2 // indirect
github.com/google/gnostic-models v0.7.0 // indirect
github.com/google/go-cmp v0.7.0 // indirect
github.com/google/uuid v1.6.0 // indirect
github.com/josharian/intern v1.0.0 // indirect
github.com/json-iterator/go v1.1.12 // indirect
github.com/mailru/easyjson v0.7.7 // indirect
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
github.com/modern-go/reflect2 v1.0.3-0.20250322232337-35a7c28c31ee // indirect
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
github.com/prometheus/client_golang v1.23.2 // indirect
github.com/prometheus/client_model v0.6.2 // indirect
github.com/prometheus/common v0.66.1 // indirect
github.com/prometheus/procfs v0.16.1 // indirect
github.com/spf13/pflag v1.0.9 // indirect
github.com/x448/float16 v0.8.4 // indirect
go.yaml.in/yaml/v2 v2.4.3 // indirect
go.yaml.in/yaml/v3 v3.0.4 // indirect
golang.org/x/net v0.47.0 // indirect
golang.org/x/oauth2 v0.30.0 // indirect
golang.org/x/sys v0.38.0 // indirect
golang.org/x/term v0.37.0 // indirect
golang.org/x/text v0.31.0 // indirect
golang.org/x/time v0.9.0 // indirect
google.golang.org/protobuf v1.36.8 // indirect
gopkg.in/evanphx/json-patch.v4 v4.13.0 // indirect
gopkg.in/inf.v0 v0.9.1 // indirect
gopkg.in/yaml.v3 v3.0.1 // indirect
k8s.io/apiextensions-apiserver v0.34.3 // indirect
k8s.io/klog/v2 v2.130.1 // indirect
k8s.io/kube-openapi v0.0.0-20250910181357-589584f1c912 // indirect
sigs.k8s.io/json v0.0.0-20250730193827-2d320260d730 // indirect
sigs.k8s.io/randfill v1.0.0 // indirect
sigs.k8s.io/structured-merge-diff/v6 v6.3.2-0.20260122202528-d9cc6641c482 // indirect
sigs.k8s.io/yaml v1.6.0 // indirect
)

View File

@ -1,571 +0,0 @@
//go:build integration
package main
import (
"context"
"fmt"
"os"
"sync"
"testing"
"time"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/client-go/kubernetes"
"k8s.io/client-go/rest"
"k8s.io/client-go/tools/leaderelection"
"k8s.io/client-go/tools/leaderelection/resourcelock"
"sigs.k8s.io/controller-runtime/pkg/envtest"
)
var (
testEnv *envtest.Environment
testConfig *rest.Config
)
type electionStateNoRelease struct {
cancel context.CancelFunc
elected chan struct{}
lost chan struct{}
}
func TestMain(m *testing.M) {
testEnv = &envtest.Environment{}
var err error
testConfig, err = testEnv.Start()
if err != nil {
fmt.Fprintf(os.Stderr, "failed to start envtest: %v\n", err)
os.Exit(1)
}
// Set up global client for the wrapper
client, err := kubernetes.NewForConfig(testConfig)
if err != nil {
fmt.Fprintf(os.Stderr, "failed to create clientset: %v\n", err)
testEnv.Stop()
os.Exit(1)
}
clientMutex.Lock()
globalClient = client
clientMutex.Unlock()
code := m.Run()
testEnv.Stop()
os.Exit(code)
}
func runElectionWithoutRelease(namespace, leaseName, identity string,
leaseDurationSec, renewDeadlineSec, retryPeriodSec int) (*electionStateNoRelease, error) {
if err := ensureClientInitialized(); err != nil {
return nil, err
}
ctx, cancel := context.WithCancel(context.Background())
state := &electionStateNoRelease{
cancel: cancel,
elected: make(chan struct{}),
lost: make(chan struct{}),
}
lock := &resourcelock.LeaseLock{
LeaseMeta: metav1.ObjectMeta{
Name: leaseName,
Namespace: namespace,
},
Client: globalClient.CoordinationV1(),
LockConfig: resourcelock.ResourceLockConfig{
Identity: identity,
},
}
le, err := leaderelection.NewLeaderElector(leaderelection.LeaderElectionConfig{
Lock: lock,
LeaseDuration: time.Duration(leaseDurationSec) * time.Second,
RenewDeadline: time.Duration(renewDeadlineSec) * time.Second,
RetryPeriod: time.Duration(retryPeriodSec) * time.Second,
ReleaseOnCancel: false,
Callbacks: leaderelection.LeaderCallbacks{
OnStartedLeading: func(ctx context.Context) {
close(state.elected)
<-ctx.Done()
},
OnStoppedLeading: func() {
close(state.lost)
},
},
})
if err != nil {
cancel()
return nil, fmt.Errorf("failed to create leader elector: %w", err)
}
go le.Run(ctx)
return state, nil
}
// TestSingleLeaderElection verifies a single candidate becomes leader.
func TestSingleLeaderElection(t *testing.T) {
ns := "default"
lease := "single-election-test"
identity := "node-1:8080"
err := runElection(ns, lease, identity, 5, 4, 1)
if err != nil {
t.Fatalf("runElection failed: %v", err)
}
// Wait for elected
key := electionKey(ns, lease)
electionMutex.Lock()
state := elections[key]
electionMutex.Unlock()
select {
case <-state.elected:
// success
case <-time.After(15 * time.Second):
t.Fatal("timed out waiting for election")
}
// Verify holder via getHolder
holder, transitions, err := getHolder(ns, lease)
if err != nil {
t.Fatalf("getHolder failed: %v", err)
}
if holder != identity {
t.Errorf("expected holder %q, got %q", identity, holder)
}
// First election — transitions should be 0 or 1
if transitions < 0 {
t.Errorf("expected non-negative transitions, got %d", transitions)
}
// Cancel the election
electionMutex.Lock()
state = elections[key]
electionMutex.Unlock()
state.cancel()
select {
case <-state.lost:
// success
case <-time.After(10 * time.Second):
t.Fatal("timed out waiting for election loss after cancel")
}
}
// TestLeaderEpoch verifies leaseTransitions increments across elections.
func TestLeaderEpoch(t *testing.T) {
ns := "default"
lease := "epoch-test"
// First election
err := runElection(ns, lease, "node-epoch-1:8080", 5, 4, 1)
if err != nil {
t.Fatalf("first runElection failed: %v", err)
}
key := electionKey(ns, lease)
electionMutex.Lock()
state1 := elections[key]
electionMutex.Unlock()
select {
case <-state1.elected:
case <-time.After(15 * time.Second):
t.Fatal("timed out on first election")
}
_, trans1, _ := getHolder(ns, lease)
// Cancel first election and wait for loss
state1.cancel()
select {
case <-state1.lost:
case <-time.After(10 * time.Second):
t.Fatal("timed out waiting for first election loss")
}
// Wait for lease to expire / be released
time.Sleep(2 * time.Second)
// Second election
err = runElection(ns, lease, "node-epoch-2:8080", 5, 4, 1)
if err != nil {
t.Fatalf("second runElection failed: %v", err)
}
electionMutex.Lock()
state2 := elections[key]
electionMutex.Unlock()
select {
case <-state2.elected:
case <-time.After(15 * time.Second):
t.Fatal("timed out on second election")
}
_, trans2, _ := getHolder(ns, lease)
if trans2 <= trans1 {
t.Errorf("expected transitions to increment: first=%d, second=%d", trans1, trans2)
}
state2.cancel()
select {
case <-state2.lost:
case <-time.After(10 * time.Second):
t.Fatal("timed out waiting for second election loss")
}
}
// TestSequentialLeadershipHandoff tests that a second candidate can acquire
// leadership after the first one releases it.
func TestSequentialLeadershipHandoff(t *testing.T) {
ns := "default"
lease := "two-candidate-test"
err1 := runElection(ns, lease, "candidate-a:8080", 5, 4, 1)
if err1 != nil {
t.Fatalf("first runElection failed: %v", err1)
}
key := electionKey(ns, lease)
electionMutex.Lock()
stateA := elections[key]
electionMutex.Unlock()
// Wait for first candidate to win
select {
case <-stateA.elected:
case <-time.After(15 * time.Second):
t.Fatal("timed out waiting for first candidate")
}
// Verify holder is candidate-a
holder, _, err := getHolder(ns, lease)
if err != nil {
t.Fatalf("getHolder failed: %v", err)
}
if holder != "candidate-a:8080" {
t.Errorf("expected candidate-a, got %q", holder)
}
// Cancel candidate-a
stateA.cancel()
select {
case <-stateA.lost:
case <-time.After(10 * time.Second):
t.Fatal("timed out waiting for candidate-a loss")
}
// Wait for lease to expire
time.Sleep(2 * time.Second)
// Start candidate-b
err2 := runElection(ns, lease, "candidate-b:8080", 5, 4, 1)
if err2 != nil {
t.Fatalf("second runElection failed: %v", err2)
}
electionMutex.Lock()
stateB := elections[key]
electionMutex.Unlock()
select {
case <-stateB.elected:
case <-time.After(15 * time.Second):
t.Fatal("timed out waiting for candidate-b")
}
holder, _, err = getHolder(ns, lease)
if err != nil {
t.Fatalf("getHolder after takeover failed: %v", err)
}
if holder != "candidate-b:8080" {
t.Errorf("expected candidate-b, got %q", holder)
}
stateB.cancel()
select {
case <-stateB.lost:
case <-time.After(10 * time.Second):
t.Fatal("timed out waiting for candidate-b loss")
}
}
// TestConcurrentCandidateElection starts two candidates simultaneously and
// verifies that exactly one wins leadership.
func TestConcurrentCandidateElection(t *testing.T) {
ns := "default"
lease := "concurrent-election-test"
type result struct {
identity string
elected bool
}
candidates := []string{"candidate-a:8080", "candidate-b:8080"}
results := make(chan result, len(candidates))
lock := func(identity string) *resourcelock.LeaseLock {
return &resourcelock.LeaseLock{
LeaseMeta: metav1.ObjectMeta{
Name: lease,
Namespace: ns,
},
Client: globalClient.CoordinationV1(),
LockConfig: resourcelock.ResourceLockConfig{
Identity: identity,
},
}
}
var wg sync.WaitGroup
for _, id := range candidates {
wg.Add(1)
go func(identity string) {
defer wg.Done()
// Short timeout: enough for one to acquire, but the loser
// times out before the winner's lease could expire.
ctx, cancel := context.WithTimeout(context.Background(), 8*time.Second)
defer cancel()
elected := make(chan struct{})
le, err := leaderelection.NewLeaderElector(leaderelection.LeaderElectionConfig{
Lock: lock(identity),
LeaseDuration: 5 * time.Second,
RenewDeadline: 3 * time.Second,
RetryPeriod: 1 * time.Second,
ReleaseOnCancel: true,
Callbacks: leaderelection.LeaderCallbacks{
OnStartedLeading: func(ctx context.Context) {
close(elected)
<-ctx.Done()
},
OnStoppedLeading: func() {},
},
})
if err != nil {
t.Errorf("NewLeaderElector(%s): %v", identity, err)
return
}
go le.Run(ctx)
select {
case <-elected:
results <- result{identity, true}
// Keep holding until context expires (8s total).
// Winner does NOT release early, so loser cannot
// re-acquire within its own 8s window.
<-ctx.Done()
case <-ctx.Done():
results <- result{identity, false}
}
}(id)
}
wg.Wait()
close(results)
winners := 0
for r := range results {
if r.elected {
winners++
t.Logf("winner: %s", r.identity)
}
}
if winners != 1 {
t.Fatalf("expected exactly 1 winner, got %d", winners)
}
}
// TestCancelElection tests that cancelling an election makes WaitLost return.
func TestCancelElection(t *testing.T) {
ns := "default"
lease := "cancel-test"
err := runElection(ns, lease, "cancel-node:8080", 5, 4, 1)
if err != nil {
t.Fatalf("runElection failed: %v", err)
}
key := electionKey(ns, lease)
electionMutex.Lock()
state := elections[key]
electionMutex.Unlock()
// Wait for elected
select {
case <-state.elected:
case <-time.After(15 * time.Second):
t.Fatal("timed out waiting for election")
}
// Cancel
state.cancel()
// WaitLost should return promptly
select {
case <-state.lost:
// success
case <-time.After(10 * time.Second):
t.Fatal("WaitLost did not return after cancel")
}
}
// TestGetHolderDuringElection verifies getHolder works while election is active.
func TestGetHolderDuringElection(t *testing.T) {
ns := "default"
lease := "active-get-holder-test"
identity := "active-node:8080"
err := runElection(ns, lease, identity, 5, 4, 1)
if err != nil {
t.Fatalf("runElection failed: %v", err)
}
key := electionKey(ns, lease)
electionMutex.Lock()
state := elections[key]
electionMutex.Unlock()
select {
case <-state.elected:
case <-time.After(15 * time.Second):
t.Fatal("timed out waiting for election")
}
// Concurrent getHolder calls during active election
var wg sync.WaitGroup
for i := 0; i < 5; i++ {
wg.Add(1)
go func() {
defer wg.Done()
holder, _, err := getHolder(ns, lease)
if err != nil {
t.Errorf("getHolder during election failed: %v", err)
return
}
if holder != identity {
t.Errorf("expected %q, got %q", identity, holder)
}
}()
}
wg.Wait()
state.cancel()
<-state.lost
}
// TestGetHolderReturnsEmptyAfterLeaderDeath verifies that after a leader stops
// renewing its lease without releasing it, getHolder returns an empty holder
// once the lease expires. This is the integration-level counterpart to the
// unit test TestGetHolderReturnsEmptyForExpiredLease.
func TestGetHolderReturnsEmptyAfterLeaderDeath(t *testing.T) {
ns := "default"
lease := "expired-leader-test"
identity := "doomed-leader:8080"
// Acquire leadership without ReleaseOnCancel so canceling simulates a dead
// leader that stops renewing and leaves the old holder until expiry.
state, err := runElectionWithoutRelease(ns, lease, identity, 5, 4, 1)
if err != nil {
t.Fatalf("runElection failed: %v", err)
}
select {
case <-state.elected:
case <-time.After(15 * time.Second):
t.Fatal("timed out waiting for election")
}
// Verify holder while active.
holder, _, err := getHolder(ns, lease)
if err != nil {
t.Fatalf("getHolder (active) failed: %v", err)
}
if holder != identity {
t.Fatalf("expected active holder %q, got %q", identity, holder)
}
// Simulate leader death: stop renewing without explicitly releasing.
state.cancel()
select {
case <-state.lost:
case <-time.After(10 * time.Second):
t.Fatal("timed out waiting for loss")
}
// Wait for the lease to expire (leaseDuration=5s, add margin).
time.Sleep(7 * time.Second)
// After expiry, getHolder must return empty holder so that the
// supervisor will attempt acquisition.
holder, _, err = getHolder(ns, lease)
if err != nil {
t.Fatalf("getHolder (expired) failed: %v", err)
}
if holder != "" {
t.Errorf("expected empty holder after lease expiry, got %q", holder)
}
}
// TestFailoverAfterLeaderDeath verifies that a new candidate can acquire
// leadership after the previous leader dies and its lease expires.
func TestFailoverAfterLeaderDeath(t *testing.T) {
ns := "default"
lease := "failover-test"
// First leader acquires without ReleaseOnCancel so canceling leaves the
// old holder in place until the lease naturally expires.
state1, err := runElectionWithoutRelease(ns, lease, "leader-1:8080", 5, 4, 1)
if err != nil {
t.Fatalf("first runElection failed: %v", err)
}
select {
case <-state1.elected:
case <-time.After(15 * time.Second):
t.Fatal("timed out waiting for first election")
}
// Simulate crash: cancel without release, wait for expiry.
state1.cancel()
<-state1.lost
time.Sleep(7 * time.Second)
// Second candidate should be able to acquire.
err = runElection(ns, lease, "leader-2:8080", 5, 4, 1)
if err != nil {
t.Fatalf("second runElection failed: %v", err)
}
key := electionKey(ns, lease)
electionMutex.Lock()
state2 := elections[key]
electionMutex.Unlock()
select {
case <-state2.elected:
// success — failover worked
case <-time.After(15 * time.Second):
t.Fatal("second candidate failed to acquire after leader death")
}
holder, _, err := getHolder(ns, lease)
if err != nil {
t.Fatalf("getHolder after failover failed: %v", err)
}
if holder != "leader-2:8080" {
t.Errorf("expected new leader %q, got %q", "leader-2:8080", holder)
}
state2.cancel()
<-state2.lost
}

View File

@ -1,489 +0,0 @@
package main
/*
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
// Trampoline to invoke C/C++ callback safely from Go via cgo.
typedef void (*holder_change_cb_t)(void* ctx,
const char* holder, size_t holderSize,
int64_t leaseTransitions);
static inline void call_holder_change_cb(holder_change_cb_t func, void* ctx,
const char* holder, size_t holderSize,
int64_t leaseTransitions) {
func(ctx, holder, holderSize, leaseTransitions);
}
*/
import "C"
import (
"context"
"fmt"
"os"
"sync"
"time"
"unsafe"
coordinationv1 "k8s.io/api/coordination/v1"
apierrors "k8s.io/apimachinery/pkg/api/errors"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/apimachinery/pkg/watch"
"k8s.io/client-go/kubernetes"
"k8s.io/client-go/rest"
"k8s.io/client-go/tools/clientcmd"
"k8s.io/client-go/tools/leaderelection"
"k8s.io/client-go/tools/leaderelection/resourcelock"
)
// electionState holds the runtime state for a single leader election.
type electionState struct {
cancel context.CancelFunc
elected chan struct{} // closed when OnStartedLeading fires
lost chan struct{} // closed when OnStoppedLeading fires
err error // set before lost is closed, if any
transitions int64 // set before elected is closed
}
// watchState holds the runtime state for a single Lease watch.
type watchState struct {
cancel context.CancelFunc
}
var (
globalClient kubernetes.Interface
clientMutex sync.Mutex
initClientFn = initClient
elections = make(map[string]*electionState)
electionMutex sync.Mutex
watches = make(map[string]*watchState)
watchMutex sync.Mutex
)
func electionKey(namespace, leaseName string) string {
return namespace + "/" + leaseName
}
func ensureClientInitialized() error {
clientMutex.Lock()
initialized := globalClient != nil
clientMutex.Unlock()
if initialized {
return nil
}
return initClientFn()
}
// initClient creates the K8s clientset from in-cluster config or KUBECONFIG.
func initClient() error {
clientMutex.Lock()
defer clientMutex.Unlock()
if globalClient != nil {
return nil
}
config, err := rest.InClusterConfig()
if err != nil {
// Fall back to KUBECONFIG
kubeconfig := os.Getenv("KUBECONFIG")
if kubeconfig == "" {
home := os.Getenv("HOME")
if home != "" {
kubeconfig = home + "/.kube/config"
}
}
config, err = clientcmd.BuildConfigFromFlags("", kubeconfig)
if err != nil {
return fmt.Errorf("failed to build k8s config: %w", err)
}
}
client, err := kubernetes.NewForConfig(config)
if err != nil {
return fmt.Errorf("failed to create k8s clientset: %w", err)
}
globalClient = client
return nil
}
// runElection starts a leader election goroutine for the given namespace/leaseName.
func runElection(namespace, leaseName, identity string,
leaseDurationSec, renewDeadlineSec, retryPeriodSec int) error {
if err := ensureClientInitialized(); err != nil {
return err
}
key := electionKey(namespace, leaseName)
electionMutex.Lock()
if _, exists := elections[key]; exists {
electionMutex.Unlock()
return fmt.Errorf("election already running for %s", key)
}
ctx, cancel := context.WithCancel(context.Background())
state := &electionState{
cancel: cancel,
elected: make(chan struct{}),
lost: make(chan struct{}),
}
elections[key] = state
electionMutex.Unlock()
lock := &resourcelock.LeaseLock{
LeaseMeta: metav1.ObjectMeta{
Name: leaseName,
Namespace: namespace,
},
Client: globalClient.CoordinationV1(),
LockConfig: resourcelock.ResourceLockConfig{
Identity: identity,
},
}
le, err := leaderelection.NewLeaderElector(leaderelection.LeaderElectionConfig{
Lock: lock,
LeaseDuration: time.Duration(leaseDurationSec) * time.Second,
RenewDeadline: time.Duration(renewDeadlineSec) * time.Second,
RetryPeriod: time.Duration(retryPeriodSec) * time.Second,
ReleaseOnCancel: true,
Callbacks: leaderelection.LeaderCallbacks{
OnStartedLeading: func(ctx context.Context) {
_, transitions, err := getHolder(namespace, leaseName)
if err == nil {
state.transitions = transitions
}
close(state.elected)
// Block until context is cancelled (leadership lost or explicit cancel)
<-ctx.Done()
},
OnStoppedLeading: func() {
close(state.lost)
// Auto-cleanup: remove from map so the same key can be reused.
electionMutex.Lock()
if elections[key] == state {
delete(elections, key)
}
electionMutex.Unlock()
},
},
})
if err != nil {
electionMutex.Lock()
delete(elections, key)
electionMutex.Unlock()
cancel()
return fmt.Errorf("failed to create leader elector: %w", err)
}
go le.Run(ctx)
return nil
}
// getHolder reads the current Lease holder identity and transitions.
func getHolder(namespace, leaseName string) (string, int64, error) {
if err := ensureClientInitialized(); err != nil {
return "", 0, err
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
lease, err := globalClient.CoordinationV1().Leases(namespace).Get(ctx, leaseName, metav1.GetOptions{})
if err != nil {
return "", 0, fmt.Errorf("failed to get lease: %w", err)
}
holder := ""
if lease.Spec.HolderIdentity != nil {
holder = *lease.Spec.HolderIdentity
}
transitions := int64(0)
if lease.Spec.LeaseTransitions != nil {
transitions = int64(*lease.Spec.LeaseTransitions)
}
// Treat expired leases as having no holder so that the C++ supervisor
// will attempt acquisition instead of going to standby.
if holder != "" && lease.Spec.RenewTime != nil && lease.Spec.LeaseDurationSeconds != nil {
expiry := lease.Spec.RenewTime.Time.Add(time.Duration(*lease.Spec.LeaseDurationSeconds) * time.Second)
if time.Now().After(expiry) {
holder = ""
}
}
return holder, transitions, nil
}
//export K8sLeaseInit
func K8sLeaseInit(errMsg **C.char) C.int {
if err := ensureClientInitialized(); err != nil {
*errMsg = C.CString(err.Error())
return -1
}
return 0
}
//export K8sLeaseRunElection
func K8sLeaseRunElection(
ns, leaseName, identity *C.char,
leaseDurationSec, renewDeadlineSec, retryPeriodSec C.int,
errMsg **C.char,
) C.int {
nsStr := C.GoString(ns)
ln := C.GoString(leaseName)
id := C.GoString(identity)
err := runElection(nsStr, ln, id,
int(leaseDurationSec), int(renewDeadlineSec), int(retryPeriodSec))
if err != nil {
*errMsg = C.CString(err.Error())
return -1
}
return 0
}
//export K8sLeaseWaitElected
func K8sLeaseWaitElected(
ns, leaseName *C.char,
timeoutSec C.int,
leaseTransitions *C.longlong,
errMsg **C.char,
) C.int {
key := electionKey(C.GoString(ns), C.GoString(leaseName))
electionMutex.Lock()
state, exists := elections[key]
electionMutex.Unlock()
if !exists {
*errMsg = C.CString("no election running for " + key)
return -1
}
timeout := time.Duration(timeoutSec) * time.Second
// Wait for elected, lost, or timeout
select {
case <-state.elected:
*leaseTransitions = C.longlong(state.transitions)
return 0
case <-state.lost:
*errMsg = C.CString("election lost before becoming leader")
return -1
case <-time.After(timeout):
state.cancel()
<-state.lost
*errMsg = C.CString("election timed out after " + fmt.Sprintf("%d", int(timeoutSec)) + "s")
return -1
}
}
//export K8sLeaseWaitLost
func K8sLeaseWaitLost(
ns, leaseName *C.char,
errMsg **C.char,
) C.int {
key := electionKey(C.GoString(ns), C.GoString(leaseName))
electionMutex.Lock()
state, exists := elections[key]
electionMutex.Unlock()
if !exists {
// Already cleaned up by OnStoppedLeading — election is over.
return 0
}
<-state.lost
if state.err != nil {
*errMsg = C.CString(state.err.Error())
return -1
}
return 0
}
//export K8sLeaseCancelElection
func K8sLeaseCancelElection(
ns, leaseName *C.char,
errMsg **C.char,
) C.int {
key := electionKey(C.GoString(ns), C.GoString(leaseName))
electionMutex.Lock()
state, exists := elections[key]
electionMutex.Unlock()
if !exists {
// Idempotent — no error if no election
return 0
}
state.cancel()
return 0
}
//export K8sLeaseGetHolder
func K8sLeaseGetHolder(
ns, leaseName *C.char,
holderIdentity **C.char,
leaseTransitions *C.longlong,
errMsg **C.char,
) C.int {
nsStr := C.GoString(ns)
ln := C.GoString(leaseName)
holder, transitions, err := getHolder(nsStr, ln)
if err != nil {
if apierrors.IsNotFound(err) {
*holderIdentity = nil
*leaseTransitions = 0
return 1
}
errStr := err.Error()
*errMsg = C.CString(errStr)
return -1
}
if holder == "" {
*holderIdentity = nil
} else {
*holderIdentity = C.CString(holder)
}
*leaseTransitions = C.longlong(transitions)
return 0
}
//export K8sLeaseWatchHolder
func K8sLeaseWatchHolder(
ns, leaseName *C.char,
callbackCtx unsafe.Pointer,
callbackFunc C.holder_change_cb_t,
errMsg **C.char,
) C.int {
nsStr := C.GoString(ns)
ln := C.GoString(leaseName)
key := electionKey(nsStr, ln)
if callbackFunc == nil {
*errMsg = C.CString("callback function is nil")
return -1
}
if err := ensureClientInitialized(); err != nil {
*errMsg = C.CString(err.Error())
return -1
}
watchMutex.Lock()
if _, exists := watches[key]; exists {
watchMutex.Unlock()
*errMsg = C.CString("watch already running for " + key)
return -1
}
ctx, cancel := context.WithCancel(context.Background())
watches[key] = &watchState{cancel: cancel}
watchMutex.Unlock()
go func() {
defer func() {
watchMutex.Lock()
delete(watches, key)
watchMutex.Unlock()
}()
for {
select {
case <-ctx.Done():
return
default:
}
watcher, err := globalClient.CoordinationV1().Leases(nsStr).Watch(ctx, metav1.ListOptions{
FieldSelector: "metadata.name=" + ln,
})
if err != nil {
select {
case <-ctx.Done():
return
default:
time.Sleep(time.Second)
continue
}
}
for event := range watcher.ResultChan() {
select {
case <-ctx.Done():
watcher.Stop()
return
default:
}
if event.Type == watch.Modified || event.Type == watch.Added {
lease, ok := event.Object.(*coordinationv1.Lease)
if !ok {
continue
}
holder := ""
if lease.Spec.HolderIdentity != nil {
holder = *lease.Spec.HolderIdentity
}
transitions := int64(0)
if lease.Spec.LeaseTransitions != nil {
transitions = int64(*lease.Spec.LeaseTransitions)
}
var holderPtr *C.char
var holderSize C.size_t
if holder != "" {
holderPtr = C.CString(holder)
holderSize = C.size_t(len(holder))
}
C.call_holder_change_cb(callbackFunc, callbackCtx,
holderPtr, holderSize, C.int64_t(transitions))
if holderPtr != nil {
C.free(unsafe.Pointer(holderPtr))
}
}
}
// Watch channel closed — retry unless cancelled
select {
case <-ctx.Done():
return
default:
time.Sleep(time.Second)
}
}
}()
return 0
}
//export K8sLeaseCancelWatch
func K8sLeaseCancelWatch(
ns, leaseName *C.char,
errMsg **C.char,
) C.int {
key := electionKey(C.GoString(ns), C.GoString(leaseName))
watchMutex.Lock()
state, exists := watches[key]
watchMutex.Unlock()
if !exists {
// Idempotent
return 0
}
state.cancel()
return 0
}
func main() {}

View File

@ -1,293 +0,0 @@
package main
import (
"context"
"fmt"
"testing"
"time"
coordinationv1 "k8s.io/api/coordination/v1"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/client-go/kubernetes"
"k8s.io/client-go/kubernetes/fake"
"k8s.io/utils/ptr"
)
// swapClient replaces globalClient and returns the old one.
func swapClient(newClient kubernetes.Interface) kubernetes.Interface {
clientMutex.Lock()
defer clientMutex.Unlock()
old := globalClient
globalClient = newClient
return old
}
// TestGetHolderWithFakeClient tests getHolder using a fake K8s clientset.
func TestGetHolderWithFakeClient(t *testing.T) {
holderID := "node-1:8080"
transitions := int32(3)
lease := &coordinationv1.Lease{
ObjectMeta: metav1.ObjectMeta{
Name: "test-lease",
Namespace: "default",
},
Spec: coordinationv1.LeaseSpec{
HolderIdentity: &holderID,
LeaseTransitions: &transitions,
},
}
fakeClient := fake.NewSimpleClientset(lease)
old := swapClient(fakeClient)
defer swapClient(old)
holder, trans, err := getHolder("default", "test-lease")
if err != nil {
t.Fatalf("getHolder failed: %v", err)
}
if holder != holderID {
t.Errorf("expected holder %q, got %q", holderID, holder)
}
if trans != int64(transitions) {
t.Errorf("expected transitions %d, got %d", transitions, trans)
}
}
// TestGetHolderNotFound tests getHolder when the Lease does not exist.
func TestGetHolderNotFound(t *testing.T) {
fakeClient := fake.NewSimpleClientset()
old := swapClient(fakeClient)
defer swapClient(old)
_, _, err := getHolder("default", "nonexistent")
if err == nil {
t.Fatal("expected error for nonexistent lease, got nil")
}
}
// TestGetHolderEmptyIdentity tests getHolder when holder is nil.
func TestGetHolderEmptyIdentity(t *testing.T) {
lease := &coordinationv1.Lease{
ObjectMeta: metav1.ObjectMeta{
Name: "empty-lease",
Namespace: "default",
},
Spec: coordinationv1.LeaseSpec{},
}
fakeClient := fake.NewSimpleClientset(lease)
old := swapClient(fakeClient)
defer swapClient(old)
holder, trans, err := getHolder("default", "empty-lease")
if err != nil {
t.Fatalf("getHolder failed: %v", err)
}
if holder != "" {
t.Errorf("expected empty holder, got %q", holder)
}
if trans != 0 {
t.Errorf("expected 0 transitions, got %d", trans)
}
}
// TestGetHolderReturnsEmptyForExpiredLease verifies that getHolder treats a
// lease whose renewTime + leaseDuration is in the past as having no holder.
// This is critical for failover: when a leader pod dies without releasing the
// lease, standbys must see an empty holder so the supervisor attempts
// acquisition instead of looping in standby.
func TestGetHolderReturnsEmptyForExpiredLease(t *testing.T) {
holderID := "dead-leader:8080"
leaseDuration := int32(5)
transitions := int32(2)
expiredRenewTime := metav1.NewMicroTime(time.Now().Add(-10 * time.Second))
lease := &coordinationv1.Lease{
ObjectMeta: metav1.ObjectMeta{
Name: "expired-lease",
Namespace: "default",
},
Spec: coordinationv1.LeaseSpec{
HolderIdentity: &holderID,
LeaseDurationSeconds: &leaseDuration,
LeaseTransitions: &transitions,
RenewTime: &expiredRenewTime,
},
}
fakeClient := fake.NewSimpleClientset(lease)
old := swapClient(fakeClient)
defer swapClient(old)
holder, trans, err := getHolder("default", "expired-lease")
if err != nil {
t.Fatalf("getHolder failed: %v", err)
}
if holder != "" {
t.Errorf("expected empty holder for expired lease, got %q", holder)
}
// Transitions should still be reported even for expired leases.
if trans != int64(transitions) {
t.Errorf("expected transitions %d, got %d", transitions, trans)
}
}
// TestGetHolderReturnsHolderForActiveLease verifies that getHolder returns the
// holder identity when the lease is still active (renewTime + leaseDuration is
// in the future).
func TestGetHolderReturnsHolderForActiveLease(t *testing.T) {
holderID := "active-leader:8080"
leaseDuration := int32(15)
transitions := int32(1)
recentRenewTime := metav1.NewMicroTime(time.Now())
lease := &coordinationv1.Lease{
ObjectMeta: metav1.ObjectMeta{
Name: "active-lease",
Namespace: "default",
},
Spec: coordinationv1.LeaseSpec{
HolderIdentity: &holderID,
LeaseDurationSeconds: &leaseDuration,
LeaseTransitions: &transitions,
RenewTime: &recentRenewTime,
},
}
fakeClient := fake.NewSimpleClientset(lease)
old := swapClient(fakeClient)
defer swapClient(old)
holder, trans, err := getHolder("default", "active-lease")
if err != nil {
t.Fatalf("getHolder failed: %v", err)
}
if holder != holderID {
t.Errorf("expected holder %q, got %q", holderID, holder)
}
if trans != int64(transitions) {
t.Errorf("expected transitions %d, got %d", transitions, trans)
}
}
// TestElectionKeyFormat tests the election key construction.
func TestElectionKeyFormat(t *testing.T) {
tests := []struct {
ns, name, want string
}{
{"default", "leader", "default/leader"},
{"kube-system", "my-lock", "kube-system/my-lock"},
{"", "bare", "/bare"},
}
for _, tc := range tests {
got := electionKey(tc.ns, tc.name)
if got != tc.want {
t.Errorf("electionKey(%q, %q) = %q, want %q", tc.ns, tc.name, got, tc.want)
}
}
}
// TestLeaseCRUDWithFakeClient tests basic Lease CRUD via the K8s API.
func TestLeaseCRUDWithFakeClient(t *testing.T) {
fakeClient := fake.NewSimpleClientset()
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
holderID := "node-a:9090"
transitions := int32(0)
lease := &coordinationv1.Lease{
ObjectMeta: metav1.ObjectMeta{
Name: "crud-test",
Namespace: "default",
},
Spec: coordinationv1.LeaseSpec{
HolderIdentity: &holderID,
LeaseTransitions: &transitions,
},
}
created, err := fakeClient.CoordinationV1().Leases("default").Create(ctx, lease, metav1.CreateOptions{})
if err != nil {
t.Fatalf("create lease failed: %v", err)
}
if *created.Spec.HolderIdentity != holderID {
t.Errorf("created holder = %q, want %q", *created.Spec.HolderIdentity, holderID)
}
newHolder := "node-b:9090"
newTransitions := int32(1)
created.Spec.HolderIdentity = &newHolder
created.Spec.LeaseTransitions = &newTransitions
updated, err := fakeClient.CoordinationV1().Leases("default").Update(ctx, created, metav1.UpdateOptions{})
if err != nil {
t.Fatalf("update lease failed: %v", err)
}
if *updated.Spec.HolderIdentity != newHolder {
t.Errorf("updated holder = %q, want %q", *updated.Spec.HolderIdentity, newHolder)
}
if *updated.Spec.LeaseTransitions != newTransitions {
t.Errorf("updated transitions = %d, want %d", *updated.Spec.LeaseTransitions, newTransitions)
}
got, err := fakeClient.CoordinationV1().Leases("default").Get(ctx, "crud-test", metav1.GetOptions{})
if err != nil {
t.Fatalf("get lease failed: %v", err)
}
if *got.Spec.HolderIdentity != newHolder {
t.Errorf("got holder = %q, want %q", *got.Spec.HolderIdentity, newHolder)
}
err = fakeClient.CoordinationV1().Leases("default").Delete(ctx, "crud-test", metav1.DeleteOptions{})
if err != nil {
t.Fatalf("delete lease failed: %v", err)
}
_, err = fakeClient.CoordinationV1().Leases("default").Get(ctx, "crud-test", metav1.GetOptions{})
if err == nil {
t.Fatal("expected error after delete, got nil")
}
}
// TestConcurrentGetHolder tests concurrent calls to getHolder.
func TestConcurrentGetHolder(t *testing.T) {
holderID := "concurrent-node:8080"
lease := &coordinationv1.Lease{
ObjectMeta: metav1.ObjectMeta{
Name: "concurrent-lease",
Namespace: "default",
},
Spec: coordinationv1.LeaseSpec{
HolderIdentity: &holderID,
LeaseTransitions: ptr.To(int32(5)),
},
}
fakeClient := fake.NewSimpleClientset(lease)
old := swapClient(fakeClient)
defer swapClient(old)
const n = 10
errCh := make(chan error, n)
for i := 0; i < n; i++ {
go func() {
holder, trans, err := getHolder("default", "concurrent-lease")
if err != nil {
errCh <- err
return
}
if holder != holderID {
errCh <- fmt.Errorf("expected holder %q, got %q", holderID, holder)
return
}
if trans != 5 {
errCh <- fmt.Errorf("expected transitions 5, got %d", trans)
return
}
errCh <- nil
}()
}
for i := 0; i < n; i++ {
if err := <-errCh; err != nil {
t.Fatalf("concurrent getHolder failed: %v", err)
}
}
}

View File

@ -1,64 +1,13 @@
find_package(yaml-cpp REQUIRED)
find_package(asio QUIET)
if(asio_FOUND)
message(STATUS "Found ASIO via find_package")
set(ASIO_INCLUDE_DIR ${asio_INCLUDE_DIR})
else()
find_path(ASIO_INCLUDE_DIR
NAMES asio.hpp
PATHS
/usr/local/include
/usr/include
${CMAKE_INSTALL_PREFIX}/include
DOC "Path to ASIO headers"
)
if(NOT ASIO_INCLUDE_DIR)
message(FATAL_ERROR "ASIO not found. Please install ASIO or set ASIO_INCLUDE_DIR manually.")
endif()
message(STATUS "Found ASIO at: ${ASIO_INCLUDE_DIR}")
endif()
set(MOONCAKE_COMMON_SOURCES
default_config.cpp
environ.cpp
)
add_library(asio_shared SHARED asio_impl.cpp)
target_compile_definitions(asio_shared
PUBLIC
ASIO_SEPARATE_COMPILATION
ASIO_DYN_LINK
)
target_include_directories(asio_shared
PUBLIC
${ASIO_INCLUDE_DIR}
)
set_target_properties(asio_shared PROPERTIES
POSITION_INDEPENDENT_CODE ON
INSTALL_RPATH "$ORIGIN"
BUILD_WITH_INSTALL_RPATH TRUE
OUTPUT_NAME "asio"
LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/mooncake-common"
)
target_link_libraries(asio_shared PUBLIC pthread)
add_library(mooncake_common
${MOONCAKE_COMMON_SOURCES}
)
target_include_directories(mooncake_common PUBLIC
$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/../include>
$<INSTALL_INTERFACE:include>
)
target_link_libraries(mooncake_common PUBLIC
yaml-cpp
jsoncpp
@ -67,5 +16,3 @@ target_link_libraries(mooncake_common PUBLIC
if (BUILD_SHARED_LIBS)
install(TARGETS mooncake_common DESTINATION lib)
endif()
install(TARGETS asio_shared DESTINATION lib)

View File

@ -1563,8 +1563,7 @@ PYBIND11_MODULE(store, m) {
const std::string &rdma_devices = "",
const std::string &master_server_addr = "127.0.0.1:50051",
const py::object &engine = py::none(),
bool enable_ssd_offload = false,
const std::string &ssd_offload_path = "") {
bool enable_ssd_offload = false) {
auto real_client = self.init_real_client();
std::shared_ptr<mooncake::TransferEngine> transfer_engine =
nullptr;
@ -1575,15 +1574,14 @@ PYBIND11_MODULE(store, m) {
return real_client->setup_real(
local_hostname, metadata_server, global_segment_size,
local_buffer_size, protocol, rdma_devices,
master_server_addr, transfer_engine, "", enable_ssd_offload,
ssd_offload_path);
master_server_addr, transfer_engine, "",
enable_ssd_offload);
},
py::arg("local_hostname"), py::arg("metadata_server"),
py::arg("global_segment_size"), py::arg("local_buffer_size"),
py::arg("protocol"), py::arg("rdma_devices"),
py::arg("master_server_addr"), py::arg("engine") = py::none(),
py::arg("enable_ssd_offload") = false,
py::arg("ssd_offload_path") = "")
py::arg("enable_ssd_offload") = false)
.def(
"setup",
[](MooncakeStorePyWrapper &self, const py::dict &config_dict) {
@ -1612,9 +1610,7 @@ PYBIND11_MODULE(store, m) {
" rdma_devices: RDMA device list.\n"
" master_server_addr: Master server address.\n"
" ipc_socket_path: IPC socket path.\n"
" enable_ssd_offload: Enable SSD offload (default false).\n"
" ssd_offload_path: SSD storage directory path (overrides env "
"var).")
" enable_ssd_offload: Enable SSD offload (default false).")
.def(
"setup_dummy",
[](MooncakeStorePyWrapper &self, size_t mem_pool_size,
@ -2271,20 +2267,6 @@ PYBIND11_MODULE(store, m) {
return self.store_->batch_get_replica_desc(keys);
},
py::arg("keys"))
.def(
"batch_replica_clear",
[](MooncakeStorePyWrapper &self,
const std::vector<std::string> &keys,
const std::string &segment_name) {
if (!self.is_client_initialized()) {
LOG(ERROR) << "Client is not initialized";
return std::vector<std::string>{};
}
py::gil_scoped_release release;
return self.store_->batch_replica_clear(keys, segment_name);
},
py::arg("keys"), py::arg("segment_name") = "",
"Clear replicas for the given keys. Requires lease to be expired.")
.def(
"create_copy_task",
[](MooncakeStorePyWrapper &self, const std::string &key,

View File

@ -33,7 +33,7 @@ fi
EXT_LDFLAGS="-L$BUILD_DIR/mooncake-transfer-engine/src"
EXT_LDFLAGS+=" -L$BUILD_DIR/mooncake-transfer-engine/src/common/base"
EXT_LDFLAGS+=" -L$BUILD_DIR/mooncake-common"
EXT_LDFLAGS+=" -L$BUILD_DIR/mooncake-asio"
EXT_LDFLAGS+=" -ltransfer_engine -lbase -lasio -lstdc++ -lnuma -lglog -libverbs -ljsoncpp"
if [ -d "/usr/local/cuda/lib64/stubs" ]; then

View File

@ -125,33 +125,17 @@ class MooncakeBackend final : public ::c10d::Backend {
}
std::string getPreferredHca(std::string location) {
static std::once_flag topo_once;
static std::shared_ptr<Topology> topology;
static TopologyMatrix matrix;
std::call_once(topo_once, [this] {
// FIXME: getLocalTopology is deprecated in TENT
topology = engine_->getLocalTopology();
if (topology) {
matrix = topology->getMatrix();
}
if (!topology || matrix.empty()) {
topology = std::make_shared<Topology>();
topology->discover();
matrix = topology->getMatrix();
}
});
auto matrix = engine_->getLocalTopology()->getMatrix();
auto it = matrix.find(location);
if (it == matrix.end()) {
LOG(INFO) << "Topology is " << topology->toJson();
LOG(INFO) << "Topology is "
<< engine_->getLocalTopology()->toJson();
LOG(ERROR) << "Topology entry not found for location: " << location;
return "";
}
if (it->second.preferred_hca.empty()) {
LOG(INFO) << "Topology is " << topology->toJson();
} else if (it->second.preferred_hca.empty()) {
LOG(INFO) << "Topology is "
<< engine_->getLocalTopology()->toJson();
LOG(ERROR) << "Preferred HCA list is empty for location: "
<< location;
return "";
}
return it->second.preferred_hca[0];
}

View File

@ -11,11 +11,9 @@
#include <transfer_engine.h>
#include <memory>
#include <atomic>
#include <mutex>
#include <thread>
#include <unordered_map>
#include <vector>
namespace mooncake {
@ -53,7 +51,6 @@ __global__ struct Task {
size_t tensorSize; // In bytes
int64_t broadcastRoot;
int bufferOffset;
uint64_t submitSequence = 0;
BatchID batchID;
void* transferGroupMeta;
};
@ -67,12 +64,6 @@ void launchReduceCpu(at::Tensor dst, size_t pos, size_t realSize, void* src,
void preloadReduceKernels();
class ConnectionContext;
struct CudaTaskSubmissionToken {
size_t task_id;
uint64_t sequence;
};
class MooncakeWorker {
public:
explicit MooncakeWorker(int cuda_device_index = -1);
@ -91,11 +82,11 @@ class MooncakeWorker {
c10d::OpType opType, size_t tensorSize, int64_t broadcastRoot,
const std::shared_ptr<TransferGroupMeta>& meta,
const std::shared_ptr<ConnectionContext>& connection_ctx,
const at::cuda::CUDAStream& issue_stream,
const std::function<void(void* dst, size_t pos, size_t realSize,
const at::cuda::CUDAStream&)>& tensorToBuffer,
const std::function<void(void* src, size_t pos, size_t realSize,
const at::cuda::CUDAStream&)>& bufferToTensor);
const at::cuda::CUDAStream& stream,
const std::function<void(void* dst, size_t pos, size_t realSize)>&
tensorToBuffer,
const std::function<void(void* src, size_t pos, size_t realSize)>&
bufferToTensor);
void Start();
@ -113,10 +104,6 @@ class MooncakeWorker {
*/
bool drainTasks(const TransferGroupMeta* meta) const;
bool waitUntilTasksSubmitted(
const std::vector<CudaTaskSubmissionToken>& tasks,
std::chrono::milliseconds timeout) const;
private:
void startWorker();
@ -135,8 +122,6 @@ class MooncakeWorker {
int cpuTaskCount = 0;
int cudaTaskCount = 0;
std::atomic<uint64_t> next_cuda_task_sequence_{1};
std::atomic<uint64_t> submitted_task_sequence_[kNumTasks_]{};
std::thread worker_thread_;
};

View File

@ -103,18 +103,17 @@ MooncakeBackend::MooncakeBackend(
const int size = distBackendOpts.group_size;
const auto& globalRanks = distBackendOpts.global_ranks_in_group;
// Memory location for device specific buffers
// always kWildcardLocation for cpu backend
std::string location = kWildcardLocation;
if (!isCpu) {
int deviceCount = 0;
cudaError_t err = cudaGetDeviceCount(&deviceCount);
if (err == cudaSuccess && deviceCount != 0) {
int deviceId_;
err = cudaGetDevice(&deviceId_);
TORCH_CHECK(!err, c10::str("Failed to get device id"));
location = GPU_PREFIX + std::to_string(deviceId_);
}
// Get device data
std::string location;
int deviceCount = 0;
cudaError_t err = cudaGetDeviceCount(&deviceCount);
if (err != cudaSuccess || deviceCount == 0) {
location = kWildcardLocation;
} else {
int deviceId_;
err = cudaGetDevice(&deviceId_);
TORCH_CHECK(!err, c10::str("Failed to get device id"));
location = GPU_PREFIX + std::to_string(deviceId_);
}
// Initialize transfer engine
@ -391,18 +390,15 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::broadcast(
return worker_->putTaskCuda(
c10d::OpType::BROADCAST, tensorSize, root, meta_, connection_ctx_,
stream,
[=](void* dst, size_t pos, size_t realSize,
const at::cuda::CUDAStream& enq_stream) {
[=](void* dst, size_t pos, size_t realSize) {
if (isRoot) {
cudaMemcpyAsync(dst, (char*)tensor.data_ptr() + pos,
realSize, cudaMemcpyDeviceToDevice,
enq_stream);
realSize, cudaMemcpyDeviceToDevice, stream);
}
},
[=](void* src, size_t pos, size_t realSize,
const at::cuda::CUDAStream& enq_stream) {
[=](void* src, size_t pos, size_t realSize) {
cudaMemcpyAsync((char*)tensor.data_ptr() + pos, src, realSize,
cudaMemcpyDeviceToDevice, enq_stream);
cudaMemcpyDeviceToDevice, stream);
});
}
}
@ -430,18 +426,16 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::allreduce(
return worker_->putTaskCuda(
c10d::OpType::ALLREDUCE, tensorSize, 0, meta_, connection_ctx_,
stream,
[=](void* dst, size_t pos, size_t realSize,
const at::cuda::CUDAStream& enq_stream) {
[=](void* dst, size_t pos, size_t realSize) {
cudaMemcpyAsync(dst, (char*)tensor.data_ptr() + pos, realSize,
cudaMemcpyDeviceToDevice, enq_stream);
cudaMemcpyDeviceToDevice, stream);
},
[=, this](void* src, size_t pos, size_t realSize,
const at::cuda::CUDAStream& enq_stream) {
[=, this](void* src, size_t pos, size_t realSize) {
cudaMemsetAsync((char*)tensor.data_ptr() + pos, 0, realSize,
enq_stream);
stream);
launchReduceKernel(tensor, pos, realSize, src, meta_->size,
opts.reduceOp, meta_->activeRanksDevice,
enq_stream);
stream);
});
}
}
@ -472,17 +466,15 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::allgather(
return worker_->putTaskCuda(
c10d::OpType::ALLGATHER, tensorSize, 0, meta_, connection_ctx_,
stream,
[=](void* dst, size_t pos, size_t realSize,
const at::cuda::CUDAStream& enq_stream) {
[=](void* dst, size_t pos, size_t realSize) {
cudaMemcpyAsync(dst, (char*)inputTensor.data_ptr() + pos,
realSize, cudaMemcpyDeviceToDevice, enq_stream);
realSize, cudaMemcpyDeviceToDevice, stream);
},
[=](void* src, size_t pos, size_t realSize,
const at::cuda::CUDAStream& enq_stream) {
[=](void* src, size_t pos, size_t realSize) {
for (const auto j : c10::irange(outputTensors_.size())) {
cudaMemcpyAsync((char*)outputTensors_[j].data_ptr() + pos,
(char*)src + j * realSize, realSize,
cudaMemcpyDeviceToDevice, enq_stream);
cudaMemcpyDeviceToDevice, stream);
}
});
}
@ -513,18 +505,16 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::_allgather_base(
return worker_->putTaskCuda(
c10d::OpType::_ALLGATHER_BASE, tensorSize, 0, meta_,
connection_ctx_, stream,
[=](void* dst, size_t pos, size_t realSize,
const at::cuda::CUDAStream& enq_stream) {
[=](void* dst, size_t pos, size_t realSize) {
cudaMemcpyAsync(dst, (char*)inputBuffer.data_ptr() + pos,
realSize, cudaMemcpyDeviceToDevice, enq_stream);
realSize, cudaMemcpyDeviceToDevice, stream);
},
[=, this](void* src, size_t pos, size_t realSize,
const at::cuda::CUDAStream& enq_stream) {
[=, this](void* src, size_t pos, size_t realSize) {
for (const auto j : c10::irange(meta_->size)) {
cudaMemcpyAsync(
(char*)outputBuffer.data_ptr() + j * tensorSize + pos,
(char*)src + j * realSize, realSize,
cudaMemcpyDeviceToDevice, enq_stream);
cudaMemcpyDeviceToDevice, stream);
}
});
}
@ -557,22 +547,20 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::_reduce_scatter_base(
return worker_->putTaskCuda(
c10d::OpType::_REDUCE_SCATTER_BASE, tensorSize, 0, meta_,
connection_ctx_, stream,
[=, this](void* dst, size_t pos, size_t realSize,
const at::cuda::CUDAStream& enq_stream) {
[=, this](void* dst, size_t pos, size_t realSize) {
for (const auto j : c10::irange(meta_->size)) {
cudaMemcpyAsync(
(char*)dst + j * realSize,
(char*)inputBuffer.data_ptr() + j * tensorSize + pos,
realSize, cudaMemcpyDeviceToDevice, enq_stream);
realSize, cudaMemcpyDeviceToDevice, stream);
}
},
[=, this](void* src, size_t pos, size_t realSize,
const at::cuda::CUDAStream& enq_stream) {
[=, this](void* src, size_t pos, size_t realSize) {
cudaMemsetAsync((char*)outputBuffer.data_ptr() + pos, 0,
realSize, enq_stream);
realSize, stream);
launchReduceKernel(outputBuffer, pos, realSize, src,
meta_->size, opts.reduceOp,
meta_->activeRanksDevice, enq_stream);
meta_->activeRanksDevice, stream);
});
}
}
@ -603,21 +591,18 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::alltoall(
return worker_->putTaskCuda(
c10d::OpType::ALLTOALL, tensorSize, 0, meta_, connection_ctx_,
stream,
[=](void* dst, size_t pos, size_t realSize,
const at::cuda::CUDAStream& enq_stream) {
[=](void* dst, size_t pos, size_t realSize) {
for (const auto j : c10::irange(inputTensors.size())) {
cudaMemcpyAsync((char*)dst + j * realSize,
(char*)inputTensors[j].data_ptr() + pos,
realSize, cudaMemcpyDeviceToDevice,
enq_stream);
realSize, cudaMemcpyDeviceToDevice, stream);
}
},
[=](void* src, size_t pos, size_t realSize,
const at::cuda::CUDAStream& enq_stream) {
[=](void* src, size_t pos, size_t realSize) {
for (const auto j : c10::irange(outputTensors.size())) {
cudaMemcpyAsync((char*)outputTensors[j].data_ptr() + pos,
(char*)src + j * realSize, realSize,
cudaMemcpyDeviceToDevice, enq_stream);
cudaMemcpyDeviceToDevice, stream);
}
});
}
@ -636,9 +621,8 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::barrier(
auto stream = at::cuda::getCurrentCUDAStream(device_index);
return worker_->putTaskCuda(
c10d::OpType::BARRIER, kBarrierDummyTensorSize, 0, meta_,
connection_ctx_, stream,
[=](void*, size_t, size_t, const at::cuda::CUDAStream&) {},
[=](void*, size_t, size_t, const at::cuda::CUDAStream&) {});
connection_ctx_, stream, [=](void*, size_t, size_t) {},
[=](void*, size_t, size_t) {});
}
}
@ -668,19 +652,17 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::reduce(
return worker_->putTaskCuda(
c10d::OpType::REDUCE, tensorSize, root, meta_, connection_ctx_,
stream,
[=](void* dst, size_t pos, size_t realSize,
const at::cuda::CUDAStream& enq_stream) {
[=](void* dst, size_t pos, size_t realSize) {
cudaMemcpyAsync(dst, (char*)tensor.data_ptr() + pos, realSize,
cudaMemcpyDeviceToDevice, enq_stream);
cudaMemcpyDeviceToDevice, stream);
},
[=, this](void* src, size_t pos, size_t realSize,
const at::cuda::CUDAStream& enq_stream) {
[=, this](void* src, size_t pos, size_t realSize) {
if (isRoot) {
cudaMemsetAsync((char*)tensor.data_ptr() + pos, 0, realSize,
enq_stream);
stream);
launchReduceKernel(tensor, pos, realSize, src, meta_->size,
opts.reduceOp, meta_->activeRanksDevice,
enq_stream);
stream);
}
});
}
@ -718,20 +700,18 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::gather(
return worker_->putTaskCuda(
c10d::OpType::GATHER, tensorSize, root, meta_, connection_ctx_,
stream,
[=](void* dst, size_t pos, size_t realSize,
const at::cuda::CUDAStream& enq_stream) {
[=](void* dst, size_t pos, size_t realSize) {
cudaMemcpyAsync(dst, (char*)inputTensor.data_ptr() + pos,
realSize, cudaMemcpyDeviceToDevice, enq_stream);
realSize, cudaMemcpyDeviceToDevice, stream);
},
[=](void* src, size_t pos, size_t realSize,
const at::cuda::CUDAStream& enq_stream) {
[=](void* src, size_t pos, size_t realSize) {
if (isRoot) {
auto outputTensors_ = outputTensors.back();
for (const auto j : c10::irange(outputTensors_.size())) {
cudaMemcpyAsync(
(char*)outputTensors_[j].data_ptr() + pos,
(char*)src + j * realSize, realSize,
cudaMemcpyDeviceToDevice, enq_stream);
cudaMemcpyDeviceToDevice, stream);
}
}
});
@ -772,22 +752,20 @@ c10::intrusive_ptr<c10d::Work> MooncakeBackend::scatter(
return worker_->putTaskCuda(
c10d::OpType::SCATTER, tensorSize, root, meta_, connection_ctx_,
stream,
[=](void* dst, size_t pos, size_t realSize,
const at::cuda::CUDAStream& enq_stream) {
[=](void* dst, size_t pos, size_t realSize) {
if (isRoot) {
auto inputTensors_ = inputTensors.back();
for (const auto j : c10::irange(inputTensors_.size())) {
cudaMemcpyAsync(
(char*)dst + j * realSize,
(char*)inputTensors_[j].data_ptr() + pos, realSize,
cudaMemcpyDeviceToDevice, enq_stream);
cudaMemcpyDeviceToDevice, stream);
}
}
},
[=](void* src, size_t pos, size_t realSize,
const at::cuda::CUDAStream& enq_stream) {
[=](void* src, size_t pos, size_t realSize) {
cudaMemcpyAsync((char*)outputTensor.data_ptr() + pos, src,
realSize, cudaMemcpyDeviceToDevice, enq_stream);
realSize, cudaMemcpyDeviceToDevice, stream);
});
}
}

View File

@ -3,9 +3,6 @@
#include <memory>
#include <thread>
#include <mooncake_worker.cuh>
#include <ATen/cuda/CUDAGraphsUtils.cuh>
#include "pg_utils.h"
namespace mooncake {
@ -33,100 +30,18 @@ class MooncakeWorkCpu : public ::c10d::Work {
class MooncakeWorkCuda : public ::c10d::Work {
public:
MooncakeWorkCuda(c10d::OpType opType, std::shared_ptr<torch::Event> event,
std::shared_ptr<TransferGroupMeta> meta,
const MooncakeWorker* worker,
std::vector<CudaTaskSubmissionToken> submitted_tasks)
: Work(-1, opType),
event_(std::move(event)),
meta_(std::move(meta)),
worker_(worker),
submitted_tasks_(std::move(submitted_tasks)) {}
std::shared_ptr<TransferGroupMeta> meta)
: Work(-1, opType), event_(std::move(event)), meta_(std::move(meta)) {}
bool isCompleted() override { return event_->query(); }
bool wait(std::chrono::milliseconds timeout) override {
// Wait until the task has been submitted to TransferEngine:
// This tries to ensure that the CUDA kernels required for the transfer
// have been launched by the time `waitUntilTasksSubmitted` returns.
//
// Why is this needed? PyTorch documentation implies that collective
// operations should be enqueued when `wait()` returns. In practice, we
// found that violating this causes hangs.
//
// Our current hypothesis for the hang is: PyTorch assumes the kernels
// needed for the transfer are already launched when `wait` returns
// true. It may then launch subsequent operations after the collective
// (e.g., `.cpu()`). Such operations may acquire a process-wide lock in
// the CUDA runtime. Also, they may rely on the data produced by the
// collective, thus causing a synchronization on enq_stream. However,
// holding that runtime lock prevents cudaMemcpy(Async) in TE/TENT from
// launching. This means the transfer can't finish, and enq_stream won't
// complete. Thus, a deadlock occurs.
// (In practice, we found that replacing all cudaMemcpyAsync in TENT
// with cuMemcpyAsync actually alleviates this, which further suggests a
// deadlock in the CUDA runtime. However, that change is too invasive
// for TE/TENT, so we do not adopt it here.)
//
// Strictly speaking, the wait is needed for another reason: The current
// stream will be blocked on the event below. Any subsequent work on
// `current_stream` will wait on that event, which effectively waits for
// the task to be done. Therefore, we must ensure all kernels needed for
// the transfer task are launched BEFORE blocking the current stream, in
// case TE/TENT use `current_stream` to launch those kernels (though it
// is rare).
//
// Please note that this logic relies on the assumption that TE/TENT
// will launch all CUDA operations in `submitTransfer`.
// Unfortunately, TcpTransport in TE and TENT currently violates this
// assumption (cudaMemcpy(Async) may be called later from a callback),
// which can cause hangs in PG when a CUDA operation such as
// `x.cpu().item()` follows the collective. For TE's TcpTransport, the
// use of cudaMemcpy on the default stream may also contribute to the
// hang.
//
// Besides, for CPU-only transports (like RdmaTransport),
// waitUntilTasksSubmitted is totally unnecessary, but we keep it for
// uniform behavior to avoid invasive changes to TE/TENT.
bool submitted = true;
if (at::cuda::currentStreamCaptureStatus() ==
c10::cuda::CaptureStatus::None) {
// Normal execution: block until tasks are submitted.
submitted =
worker_->waitUntilTasksSubmitted(submitted_tasks_, timeout);
} else {
// During CUDA graph capture, kernels are recorded but not actually
// executed. The enqueueTaskKernel would never run, so
// waitUntilTasksSubmitted would hang because the CPU worker thread
// never sees task.active == true.
//
// Note that this also means NvlinkTransport (and TcpTransport too,
// of course) won't work with CUDA Graphs: Kernels launched inside
// TE/TENT can't be captured by the graph, and during replay they
// are not ordered with the graph execution. This may trigger the
// same deadlock described above.
}
if (!submitted) return false;
// Once all tasks have been submitted, use the event to synchronize
// the current stream and the enqueue stream, but do not wait on this
// event.
//
// See PyTorch docs for more details:
// https://docs.pytorch.org/docs/stable/distributed.html#synchronous-and-asynchronous-collective-operations
// "wait() - in the case of CPU collectives, will block the process
// until the operation is completed. In the case of CUDA collectives,
// will block the currently active CUDA stream until the operation
// is completed (but will not block the CPU)."
auto current_stream = at::cuda::getCurrentCUDAStream();
event_->block(current_stream);
return true;
return true; // This should be a no-op
}
protected:
std::shared_ptr<torch::Event> event_;
std::shared_ptr<TransferGroupMeta> meta_;
const MooncakeWorker* worker_;
std::vector<CudaTaskSubmissionToken> submitted_tasks_;
};
class MooncakeBarrierWorkCuda : public MooncakeWorkCuda {
@ -134,33 +49,29 @@ class MooncakeBarrierWorkCuda : public MooncakeWorkCuda {
using MooncakeWorkCuda::MooncakeWorkCuda;
bool wait(std::chrono::milliseconds timeout) override {
// Skip host-side synchronization during CUDA graph capture.
// cudaEventSynchronize is not permitted while a stream is capturing.
if (at::cuda::currentStreamCaptureStatus() !=
c10::cuda::CaptureStatus::None) {
// We still need stream-level synchronization so that subsequent
// operations on the capture stream are ordered after the barrier
// task on the enqueue stream.
auto current_stream = at::cuda::getCurrentCUDAStream();
event_->block(current_stream);
return true;
}
if (timeout == kNoTimeout) {
event_->synchronize();
return true;
}
BackoffWaiter waiter(
BackoffWaiterConfig::constantSleep(std::chrono::microseconds(10)));
return waiter.wait_for(timeout, [this] { return event_->query(); });
auto start = std::chrono::steady_clock::now();
while (!event_->query()) {
auto now = std::chrono::steady_clock::now();
auto elapsed =
std::chrono::duration_cast<std::chrono::milliseconds>(now -
start);
if (elapsed >= timeout) {
return false;
}
std::this_thread::sleep_for(std::chrono::microseconds(10));
}
return true;
}
};
__global__ void enqueueTaskKernel(c10d::OpType opType, size_t tensorSize,
int64_t broadcastRoot, int bufferOffset,
uint64_t submitSequence, void* meta,
Task* tasks, int numRanks,
void* meta, Task* tasks, int numRanks,
const bool* activeRanks,
int* activeRanksTensor, size_t taskId) {
// Copy task into slot
@ -168,16 +79,15 @@ __global__ void enqueueTaskKernel(c10d::OpType opType, size_t tensorSize,
tasks[taskId].tensorSize = tensorSize;
tasks[taskId].broadcastRoot = broadcastRoot;
tasks[taskId].bufferOffset = bufferOffset;
tasks[taskId].submitSequence = submitSequence;
tasks[taskId].transferGroupMeta = meta;
// Publish task metadata before notifying the host worker thread.
__threadfence_system();
// Mark active
__threadfence(); // Ensure writes visible to host
tasks[taskId].active = true;
// Spin-wait until CPU proxy sets DONE
while (tasks[taskId].active) {
__threadfence_system();
__threadfence();
}
for (int i = 0; i < numRanks; ++i) {
activeRanksTensor[i] = activeRanks[i] ? 1 : 0;
@ -401,8 +311,6 @@ MooncakeWorker::MooncakeWorker(int cuda_device_index)
}
for (size_t i = 0; i < kNumTasks_; ++i) {
tasks_[i].active = false;
tasks_[i].submitSequence = 0;
submitted_task_sequence_[i].store(0, std::memory_order_relaxed);
}
}
@ -492,65 +400,45 @@ c10::intrusive_ptr<c10d::Work> MooncakeWorker::putTaskCuda(
c10d::OpType opType, size_t tensorSize, int64_t broadcastRoot,
const std::shared_ptr<TransferGroupMeta>& meta,
const std::shared_ptr<ConnectionContext>& connection_ctx,
const at::cuda::CUDAStream& issue_stream,
const std::function<void(void* dst, size_t pos, size_t realSize,
const at::cuda::CUDAStream&)>& tensorToBuffer,
const std::function<void(void* src, size_t pos, size_t realSize,
const at::cuda::CUDAStream&)>& bufferToTensor) {
const at::cuda::CUDAStream& stream,
const std::function<void(void* dst, size_t pos, size_t realSize)>&
tensorToBuffer,
const std::function<void(void* src, size_t pos, size_t realSize)>&
bufferToTensor) {
connection_ctx->waitUntilNewRanksConnected();
// TORCH_CHECK(tensorSize * meta->size < kBufferSize, "Too large!");
// Alternately use even-odd items to maintain tasks
size_t chunkSize = ((kBufferSize - 1) / meta->size) & ~(size_t)7;
// Get a non-blocking stream for enqueue:
// The incoming `issue_stream` may be the Null Stream, which enforces
// implicit synchronization semantics. Launching a spin-wait kernel
// (enqueueTaskKernel) on such a stream can introduce potential deadlock.
at::cuda::CUDAStream enq_stream =
at::cuda::getStreamFromPool(false, issue_stream.device_index());
// Synchronize: enq_stream waits for issue_stream
auto event_start = std::make_shared<torch::Event>(torch::kCUDA);
event_start->record(issue_stream);
event_start->block(enq_stream);
std::vector<CudaTaskSubmissionToken> submitted_tasks;
submitted_tasks.reserve((tensorSize + chunkSize - 1) / chunkSize);
for (size_t pos = 0; pos < tensorSize; pos += chunkSize) {
size_t realSize = min(tensorSize, pos + chunkSize) - pos;
int taskId = cudaTaskCount % 2 + 2;
int bufferOffset = meta->taskCount % 2;
const uint64_t taskSequence =
next_cuda_task_sequence_.fetch_add(1, std::memory_order_relaxed);
submitted_tasks.push_back(
{.task_id = static_cast<size_t>(taskId), .sequence = taskSequence});
tensorToBuffer(
(void*)meta->segmentInfos[meta->rank].send_buffer[bufferOffset],
pos, realSize, enq_stream);
pos, realSize);
hasCallback_[taskId] = false;
enqueueTaskKernel<<<1, 1, 0, enq_stream>>>(
opType, realSize, broadcastRoot, bufferOffset, taskSequence,
meta.get(), tasks_device_, meta->size, meta->activeRanksDevice,
enqueueTaskKernel<<<1, 1, 0, stream>>>(
opType, realSize, broadcastRoot, bufferOffset, meta.get(),
tasks_device_, meta->size, meta->activeRanksDevice,
meta->activeRanksTensor.data_ptr<int>(), taskId);
bufferToTensor(
(void*)meta->segmentInfos[meta->rank].recv_buffer[bufferOffset],
pos, realSize, enq_stream);
pos, realSize);
++cudaTaskCount;
++meta->taskCount;
}
auto event_end = std::make_shared<torch::Event>(torch::kCUDA);
event_end->record(enq_stream);
auto event = std::make_shared<torch::Event>(torch::kCUDA);
event->record(stream);
if (opType == c10d::OpType::BARRIER) {
return c10::make_intrusive<MooncakeBarrierWorkCuda>(
opType, event_end, meta, this, std::move(submitted_tasks));
return c10::make_intrusive<MooncakeBarrierWorkCuda>(opType, event,
meta);
}
return c10::make_intrusive<MooncakeWorkCuda>(opType, event_end, meta, this,
std::move(submitted_tasks));
return c10::make_intrusive<MooncakeWorkCuda>(opType, event, meta);
}
} // namespace mooncake

View File

@ -35,36 +35,6 @@ bool MooncakeWorker::drainTasks(const TransferGroupMeta* meta) const {
});
}
bool MooncakeWorker::waitUntilTasksSubmitted(
const std::vector<CudaTaskSubmissionToken>& tasks,
std::chrono::milliseconds timeout) const {
if (tasks.empty()) {
return true;
}
auto submitted = [this, &tasks] {
for (const auto& task : tasks) {
if (task.task_id >= kNumTasks_) {
LOG(ERROR) << "Invalid task id.";
return true;
}
if (submitted_task_sequence_[task.task_id].load(
std::memory_order_acquire) < task.sequence) {
return false;
}
}
return true;
};
BackoffWaiter waiter(
BackoffWaiterConfig::constantSleep(std::chrono::microseconds(10)));
if (timeout == kNoTimeout) {
waiter.wait(submitted);
return true;
}
return waiter.wait_for(timeout, submitted);
}
void MooncakeWorker::startWorker() {
running_ = true;
worker_thread_ = std::thread([this] {
@ -91,10 +61,7 @@ void MooncakeWorker::startWorker() {
group->rank != task.broadcastRoot) ||
task.opType == c10d::OpType::BARRIER;
if (task_status[i].load(std::memory_order_acquire) == IDLE) {
const auto submit_sequence = task.submitSequence;
if (skipTransfer) {
submitted_task_sequence_[i].store(
submit_sequence, std::memory_order_release);
task_status[i].store(TRANSFERRED_1,
std::memory_order_release);
continue;
@ -166,8 +133,6 @@ void MooncakeWorker::startWorker() {
task.batchID =
group->engine->allocateBatchID(entries.size());
group->engine->submitTransfer(task.batchID, entries);
submitted_task_sequence_[i].store(
submit_sequence, std::memory_order_release);
activeTime[i] = clock::now();
task_status[i].store(TRANSFERRED_1,
std::memory_order_release);

View File

@ -39,7 +39,7 @@ CGO_LDFLAGS="-L${BUILD_DIR}/mooncake-store/src"
CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-store/src/cachelib_memory_allocator"
CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-transfer-engine/src"
CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-transfer-engine/src/common/base"
CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-common"
CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-asio"
CGO_LDFLAGS+=" -lmooncake_store -lcachelib_memory_allocator -ltransfer_engine -lbase -lasio"
CGO_LDFLAGS+=" -lstdc++ -lnuma -lglog -lgflags -libverbs -ljsoncpp -lzstd -lcurl"

View File

@ -1,15 +1,8 @@
#pragma once
#include <algorithm>
#include <array>
#include <atomic>
#include <chrono>
#include <iomanip>
#include <mutex>
#include <optional>
#include <sstream>
#include <thread>
#include <unordered_set>
#include <vector>
#include <ylt/metric/counter.hpp>
#include <ylt/metric/histogram.hpp>
@ -51,56 +44,6 @@ const inline std::map<std::string, std::string> merge_labels(
return merged_labels;
}
inline std::string format_metric_rate(double value, const char* suffix) {
const double KB = 1024.0;
const double MB = KB * 1024.0;
const double GB = MB * 1024.0;
const double TB = GB * 1024.0;
std::ostringstream oss;
oss << std::fixed << std::setprecision(2);
if (value >= TB) {
oss << value / TB << " T" << suffix;
} else if (value >= GB) {
oss << value / GB << " G" << suffix;
} else if (value >= MB) {
oss << value / MB << " M" << suffix;
} else if (value >= KB) {
oss << value / KB << " K" << suffix;
} else {
oss << value << " " << suffix;
}
return oss.str();
}
inline std::string format_metric_bandwidth(uint64_t total_bytes,
double elapsed_seconds) {
return format_metric_rate(total_bytes / elapsed_seconds, "B/s");
}
inline uint64_t elapsed_us_since(
std::chrono::steady_clock::time_point start_time) {
return static_cast<uint64_t>(
std::chrono::duration_cast<std::chrono::microseconds>(
std::chrono::steady_clock::now() - start_time)
.count());
}
template <typename Result, typename Operation, typename SuccessFn,
typename ObserveFn>
Result execute_timed_operation(Operation&& operation, SuccessFn&& success_fn,
ObserveFn&& observe_fn) {
const auto start_time = std::chrono::steady_clock::now();
Result result = std::forward<Operation>(operation)();
if (std::forward<SuccessFn>(success_fn)(result)) {
std::forward<ObserveFn>(observe_fn)(elapsed_us_since(start_time),
result);
}
return result;
}
enum class TransferOperationKind { kRead, kWrite };
struct TransferMetric {
TransferMetric(std::map<std::string, std::string> labels = {})
: total_read_bytes("mooncake_transfer_read_bytes", "Total bytes read",
@ -116,8 +59,7 @@ struct TransferMetric {
get_latency_us("mooncake_transfer_get_latency",
"Get transfer latency (us)", kLatencyBucket, labels),
put_latency_us("mooncake_transfer_put_latency",
"Put transfer latency (us)", kLatencyBucket, labels),
start_time_(std::chrono::steady_clock::now()) {}
"Put transfer latency (us)", kLatencyBucket, labels) {}
ylt::metric::counter_t total_read_bytes;
ylt::metric::counter_t total_write_bytes;
@ -135,7 +77,7 @@ struct TransferMetric {
put_latency_us.serialize(str);
}
std::string summary_metrics(bool include_bandwidth = true) {
std::string summary_metrics() {
std::stringstream ss;
ss << "=== Transfer Metrics Summary ===\n";
@ -144,14 +86,6 @@ struct TransferMetric {
auto write_bytes = total_write_bytes.value();
ss << "Total Read: " << byte_size_to_string(read_bytes) << "\n";
ss << "Total Write: " << byte_size_to_string(write_bytes) << "\n";
if (include_bandwidth) {
ss << "Average Read Throughput: "
<< format_metric_bandwidth(read_bytes, elapsed_seconds())
<< "\n";
ss << "Average Write Throughput: "
<< format_metric_bandwidth(write_bytes, elapsed_seconds())
<< "\n";
}
// Latency summaries
ss << "\n=== Latency Summary (microseconds) ===\n";
@ -166,14 +100,6 @@ struct TransferMetric {
}
private:
std::chrono::steady_clock::time_point start_time_;
double elapsed_seconds() const {
const auto elapsed = std::chrono::duration<double>(
std::chrono::steady_clock::now() - start_time_);
return std::max(elapsed.count(), 1e-9);
}
std::string format_latency_summary(ylt::metric::histogram_t& hist) {
// Access the internal sum and bucket counts
auto sum_ptr =
@ -346,322 +272,9 @@ struct MasterClientMetric {
}
};
struct TransferOperationMetric {
std::array<std::string, 1> op_names = {"op_name"};
explicit TransferOperationMetric(
std::map<std::string, std::string> labels = {})
: read_op_count("mooncake_transfer_read_operation_count",
"Total read operations by interface type", labels,
op_names),
read_op_bytes("mooncake_transfer_read_operation_bytes",
"Total read bytes by interface type", labels, op_names),
read_op_latency_us("mooncake_transfer_read_operation_latency",
"Read operation latency by interface type (us)",
kLatencyBucket, labels, op_names),
write_op_count("mooncake_transfer_write_operation_count",
"Total write operations by interface type", labels,
op_names),
write_op_bytes("mooncake_transfer_write_operation_bytes",
"Total write bytes by interface type", labels,
op_names),
write_op_latency_us("mooncake_transfer_write_operation_latency",
"Write operation latency by interface type (us)",
kLatencyBucket, labels, op_names) {}
ylt::metric::hybrid_counter_1t read_op_count;
ylt::metric::hybrid_counter_1t read_op_bytes;
ylt::metric::hybrid_histogram_1t read_op_latency_us;
ylt::metric::hybrid_counter_1t write_op_count;
ylt::metric::hybrid_counter_1t write_op_bytes;
ylt::metric::hybrid_histogram_1t write_op_latency_us;
void Observe(TransferOperationKind kind, const std::string& op_name,
uint64_t bytes, uint64_t latency_us) {
const std::array<std::string, 1> label = {op_name};
{
std::lock_guard<std::mutex> lock(observed_ops_mutex_);
if (kind == TransferOperationKind::kRead) {
observed_read_ops_.insert(op_name);
} else {
observed_write_ops_.insert(op_name);
}
}
if (kind == TransferOperationKind::kRead) {
read_op_count.inc(label);
read_op_bytes.inc(label, bytes);
read_op_latency_us.observe(label, latency_us);
} else {
write_op_count.inc(label);
write_op_bytes.inc(label, bytes);
write_op_latency_us.observe(label, latency_us);
}
}
void serialize(std::string& str) {
read_op_count.serialize(str);
read_op_bytes.serialize(str);
read_op_latency_us.serialize(str);
write_op_count.serialize(str);
write_op_bytes.serialize(str);
write_op_latency_us.serialize(str);
}
std::string summary_metrics() {
std::stringstream ss;
ss << "=== Interface Operation Metrics Summary ===\n";
ss << format_operation_group_summary(
"Read Interfaces", snapshot_operations(observed_read_ops_),
read_op_count, read_op_bytes, read_op_latency_us)
<< "\n";
ss << format_operation_group_summary(
"Write Interfaces", snapshot_operations(observed_write_ops_),
write_op_count, write_op_bytes, write_op_latency_us);
return ss.str();
}
private:
std::mutex observed_ops_mutex_;
std::unordered_set<std::string> observed_read_ops_;
std::unordered_set<std::string> observed_write_ops_;
std::vector<std::string> snapshot_operations(
const std::unordered_set<std::string>& source) {
std::lock_guard<std::mutex> lock(observed_ops_mutex_);
std::vector<std::string> ops(source.begin(), source.end());
std::sort(ops.begin(), ops.end());
return ops;
}
std::string format_operation_group_summary(
const std::string& group_name, const std::vector<std::string>& ops,
ylt::metric::hybrid_counter_1t& op_count,
ylt::metric::hybrid_counter_1t& op_bytes,
ylt::metric::hybrid_histogram_1t& op_latency_us) {
std::stringstream ss;
ss << group_name << ":\n";
if (ops.empty()) {
ss << "No data";
return ss.str();
}
auto bucket_counts = op_latency_us.get_bucket_counts();
bool found_any = false;
for (const auto& op_name : ops) {
const std::array<std::string, 1> label = {op_name};
const int64_t total_count = op_count.value(label);
if (total_count == 0) {
continue;
}
found_any = true;
ss << op_name << ": count=" << total_count << ", bytes="
<< byte_size_to_string(
static_cast<uint64_t>(op_bytes.value(label)));
int64_t p95_target = (total_count * 95) / 100;
int64_t cumulative = 0;
double p95_bucket = 0;
for (size_t i = 0;
i < bucket_counts.size() && i < kLatencyBucket.size(); ++i) {
cumulative += bucket_counts[i]->value(label);
if (cumulative >= p95_target && p95_bucket == 0) {
p95_bucket = kLatencyBucket[i];
break;
}
}
if (p95_bucket > 0) {
ss << ", p95<" << p95_bucket << "μs";
}
double max_bucket = 0;
for (size_t i = bucket_counts.size(); i > 0; --i) {
const size_t idx = i - 1;
if (idx < kLatencyBucket.size() &&
bucket_counts[idx]->value(label) > 0) {
max_bucket = kLatencyBucket[idx];
break;
}
}
if (max_bucket > 0) {
ss << ", max<" << max_bucket << "μs";
}
ss << "\n";
}
if (!found_any) {
ss << "No data";
}
return ss.str();
}
};
// SSD latency bucket: microseconds, tuned for SSD/network storage
// Range: 50us (high-end NVMe) to 30s (3fs/nfs large object batch writes)
inline const std::vector<double> kSsdLatencyBucket = {
50, 100, 200, // <200us (high-end NVMe)
500, 1000, 2000, 5000, 10000, // 500us - 10ms
20000, 50000, 100000, 200000, // 10ms - 200ms
500000, 1000000, 2000000, 5000000, // 500ms - 5s
10000000, 30000000 // 10s - 30s (3fs/nfs)
};
struct SsdMetric {
SsdMetric(std::map<std::string, std::string> labels = {})
: ssd_read_bytes("mooncake_ssd_read_bytes_total",
"Total bytes read from SSD", labels),
ssd_write_bytes("mooncake_ssd_write_bytes_total",
"Total bytes written to SSD", labels),
ssd_read_ops("mooncake_ssd_read_ops_total",
"Total number of SSD read operations (key count)",
labels),
ssd_write_ops("mooncake_ssd_write_ops_total",
"Total number of SSD write operations (key count)",
labels),
ssd_read_latency_us("mooncake_ssd_read_latency_us",
"SSD BatchLoad latency per batch (us)",
kSsdLatencyBucket, labels),
ssd_write_latency_us("mooncake_ssd_write_latency_us",
"SSD BatchOffload latency per batch (us)",
kSsdLatencyBucket, labels),
ssd_total_bytes("mooncake_ssd_total_bytes_total",
"Total bytes read and written to SSD", labels),
ssd_total_ops("mooncake_ssd_total_ops_total",
"Total number of SSD operations (key count)", labels),
ssd_total_latency_us("mooncake_ssd_total_latency_us",
"SSD total latency per batch (us)",
kSsdLatencyBucket, labels),
ssd_read_latency_summary("mooncake_ssd_read_latency_summary_us",
"SSD read latency quantiles (us)",
{0.5, 0.9, 0.99}, labels),
ssd_write_latency_summary("mooncake_ssd_write_latency_summary_us",
"SSD write latency quantiles (us)",
{0.5, 0.9, 0.99}, labels),
ssd_total_latency_summary("mooncake_ssd_total_latency_summary_us",
"SSD total latency quantiles (us)",
{0.5, 0.9, 0.99}, labels),
start_time_(std::chrono::steady_clock::now()) {}
ylt::metric::counter_t ssd_read_bytes;
ylt::metric::counter_t ssd_write_bytes;
ylt::metric::counter_t ssd_read_ops;
ylt::metric::counter_t ssd_write_ops;
ylt::metric::histogram_t ssd_read_latency_us;
ylt::metric::histogram_t ssd_write_latency_us;
ylt::metric::counter_t ssd_total_bytes;
ylt::metric::counter_t ssd_total_ops;
ylt::metric::histogram_t ssd_total_latency_us;
ylt::metric::summary_t ssd_read_latency_summary;
ylt::metric::summary_t ssd_write_latency_summary;
ylt::metric::summary_t ssd_total_latency_summary;
std::chrono::steady_clock::time_point start_time_;
void serialize(std::string& str) {
ssd_read_bytes.serialize(str);
ssd_write_bytes.serialize(str);
ssd_read_ops.serialize(str);
ssd_write_ops.serialize(str);
ssd_read_latency_us.serialize(str);
ssd_write_latency_us.serialize(str);
ssd_total_bytes.serialize(str);
ssd_total_ops.serialize(str);
ssd_total_latency_us.serialize(str);
ssd_read_latency_summary.serialize(str);
ssd_write_latency_summary.serialize(str);
ssd_total_latency_summary.serialize(str);
}
std::string summary_metrics() {
std::stringstream ss;
ss << "=== SSD Metrics Summary ===" << "\n";
auto read_bytes = ssd_read_bytes.value();
auto write_bytes = ssd_write_bytes.value();
auto read_ops = ssd_read_ops.value();
auto write_ops = ssd_write_ops.value();
auto elapsed_s = std::chrono::duration<double>(
std::chrono::steady_clock::now() - start_time_)
.count();
ss << "SSD Read: " << byte_size_to_string(read_bytes)
<< ", ops=" << read_ops;
if (elapsed_s > 0 && read_bytes > 0) {
ss << ", throughput="
<< byte_size_to_string(
static_cast<int64_t>(read_bytes / elapsed_s))
<< "/s";
ss << ", IOPS=" << std::fixed << std::setprecision(1)
<< (read_ops / elapsed_s);
}
ss << "\n";
ss << "SSD Write: " << byte_size_to_string(write_bytes)
<< ", ops=" << write_ops;
if (elapsed_s > 0 && write_bytes > 0) {
ss << ", throughput="
<< byte_size_to_string(
static_cast<int64_t>(write_bytes / elapsed_s))
<< "/s";
ss << ", IOPS=" << std::fixed << std::setprecision(1)
<< (write_ops / elapsed_s);
}
ss << "\n";
auto total_bytes = ssd_total_bytes.value();
auto total_ops = ssd_total_ops.value();
ss << "SSD Total: " << byte_size_to_string(total_bytes)
<< ", ops=" << total_ops;
if (elapsed_s > 0 && total_bytes > 0) {
ss << ", throughput="
<< byte_size_to_string(
static_cast<int64_t>(total_bytes / elapsed_s))
<< "/s";
ss << ", IOPS=" << std::fixed << std::setprecision(1)
<< (total_ops / elapsed_s);
}
ss << "\n";
ss << "\n" << "=== SSD Latency Summary (microseconds) ===" << "\n";
ss << "Read: " << format_summary_percentiles(ssd_read_latency_summary)
<< "\n";
ss << "Write: " << format_summary_percentiles(ssd_write_latency_summary)
<< "\n";
ss << "Total: " << format_summary_percentiles(ssd_total_latency_summary)
<< "\n";
return ss.str();
}
private:
std::string format_summary_percentiles(ylt::metric::summary_t& summary) {
double sum = 0;
uint64_t count = 0;
auto rates = summary.get_rates(sum, count);
if (count == 0) {
return "No data";
}
std::stringstream ss;
ss << std::fixed << std::setprecision(1);
ss << "count=" << count;
if (rates.size() >= 1) ss << ", p50=" << rates[0] << "us";
if (rates.size() >= 2) ss << ", p90=" << rates[1] << "us";
if (rates.size() >= 3) ss << ", p99=" << rates[2] << "us";
if (count > 0) {
ss << ", avg=" << (sum / count) << "us";
}
return ss.str();
}
};
struct ClientMetric {
TransferMetric transfer_metric;
MasterClientMetric master_client_metric;
TransferOperationMetric transfer_operation_metric;
SsdMetric ssd_metric;
/**
* @brief Creates a ClientMetric instance based on environment variables
@ -675,44 +288,25 @@ struct ClientMetric {
* (default: 0, 0 = collect but don't report)
*/
static std::unique_ptr<ClientMetric> Create(
const std::map<std::string, std::string>& labels = {},
bool master_rpc_metrics_enabled = true);
void ObserveTransferOperation(TransferOperationKind kind,
const std::string& op_name, uint64_t bytes,
uint64_t latency_us) {
transfer_operation_metric.Observe(kind, op_name, bytes, latency_us);
}
const std::map<std::string, std::string>& labels = {});
void serialize(std::string& str);
std::string summary_metrics();
uint64_t GetReportingInterval() const { return metrics_interval_seconds_; }
explicit ClientMetric(uint64_t interval_seconds = 0,
const std::map<std::string, std::string>& labels = {},
bool bandwidth_reporting_enabled = true,
bool master_rpc_metrics_enabled = true);
explicit ClientMetric(
uint64_t interval_seconds = 0,
const std::map<std::string, std::string>& labels = {});
~ClientMetric();
private:
struct TransferSnapshot {
uint64_t read_bytes;
uint64_t write_bytes;
std::chrono::steady_clock::time_point timestamp;
};
// Metrics reporting thread management
std::jthread metrics_reporting_thread_;
std::atomic<bool> should_stop_metrics_thread_{false};
uint64_t metrics_interval_seconds_{0};
bool bandwidth_reporting_enabled_{true};
bool master_rpc_metrics_enabled_{true};
std::mutex snapshot_mutex_;
std::optional<TransferSnapshot> last_report_snapshot_;
void StartMetricsReportingThread();
void StopMetricsReportingThread();
std::string BuildBandwidthReport();
};
}; // namespace mooncake
}; // namespace mooncake

View File

@ -24,7 +24,6 @@
#include "master_metric_manager.h"
#include "count_min_sketch.h"
#include "local_hot_cache.h"
#include "pinned_buffer_pool.h"
namespace mooncake {
@ -62,8 +61,6 @@ class Client {
public:
~Client();
const UUID& getClientId() const { return client_id_; }
/**
* @brief Creates and initializes a new Client instance
* @param local_hostname Local host address (IP:Port)
@ -447,15 +444,6 @@ class Client {
return master_client_.CalcCacheStats();
}
void ObserveTransferOperation(TransferOperationKind kind,
const std::string& op_name, uint64_t bytes,
uint64_t latency_us) {
if (metrics_ != nullptr) {
metrics_->ObserveTransferOperation(kind, op_name, bytes,
latency_us);
}
}
// For Prometheus-style metrics
tl::expected<std::string, ErrorCode> SerializeMetrics() {
if (metrics_ == nullptr) {
@ -466,10 +454,6 @@ class Client {
return str;
}
SsdMetric* GetSsdMetricPtr() {
return metrics_ ? &metrics_->ssd_metric : nullptr;
}
[[nodiscard]] std::string GetTransportEndpoint() {
return transfer_engine_->getLocalIpAndPort();
}
@ -672,9 +656,6 @@ class Client {
const std::string protocol_;
// Client persistent thread pool for async operations
// Pinned host memory pool for GPU D2H staging (must outlive
// write_thread_pool_)
std::unique_ptr<PinnedBufferPool> pinned_buffer_pool_;
ThreadPool write_thread_pool_;
std::shared_ptr<StorageBackend> storage_backend_;

View File

@ -7,7 +7,6 @@
#include "pyclient.h"
#include "real_client.h"
#include "shm_helper.h"
#include "client_metric.h"
#include <memory>
namespace mooncake {
@ -26,8 +25,7 @@ class DummyClient : public PyClient {
const std::string &master_server_addr,
const std::shared_ptr<TransferEngine> &transfer_engine,
const std::string &ipc_socket_path,
bool enable_ssd_offload = false,
const std::string &ssd_offload_path = "") {
bool enable_ssd_offload = false) {
// Dummy client does not support real setup
return -1;
};
@ -151,12 +149,6 @@ class DummyClient : public PyClient {
batch_get_replica_desc(const std::vector<std::string> &keys);
std::vector<Replica::Descriptor> get_replica_desc(const std::string &key);
std::vector<std::string> batch_replica_clear(
const std::vector<std::string> &keys,
const std::string &segment_name = "") override {
return {};
}
int tearDownAll();
int health_check() override;
@ -204,22 +196,6 @@ class DummyClient : public PyClient {
[[nodiscard]] std::vector<tl::expected<ResultType, ErrorCode>>
invoke_batch_rpc(size_t input_size, Args &&...args);
template <auto ServiceMethod, typename... Args>
int invoke_observed_void_rpc(TransferOperationKind kind,
const char *op_name, size_t bytes, bool batch,
Args &&...args) {
auto result = execute_timed_operation<tl::expected<void, ErrorCode>>(
[&]() {
return invoke_rpc<ServiceMethod, void>(
std::forward<Args>(args)...);
},
[](const auto &ret) { return ret.has_value(); },
[&](uint64_t latency_us, const auto &) {
ObserveTransferMetric(kind, op_name, bytes, latency_us, batch);
});
return to_py_ret(result);
}
/**
* @brief Accessor for the coro_rpc_client pool. Since coro_rpc_client
* pool cannot reconnect to a different address, a new coro_rpc_client
@ -278,11 +254,6 @@ class DummyClient : public PyClient {
// Ascend physical device id for dummy-real RPC to real, set in setup_dummy
int32_t device_id_ = 0;
std::unique_ptr<ClientMetric> metrics_;
void ObserveTransferMetric(TransferOperationKind kind, const char *op_name,
size_t bytes, uint64_t latency_us, bool batch);
};
} // namespace mooncake

View File

@ -3,17 +3,13 @@
#include "client_service.h"
#include "client_buffer.hpp"
#include "storage_backend.h"
#include "pinned_buffer_pool.h"
namespace mooncake {
struct SsdMetric;
class FileStorage {
public:
FileStorage(const FileStorageConfig& config, std::shared_ptr<Client> client,
const std::string& local_rpc_addr,
SsdMetric* ssd_metric = nullptr);
const std::string& local_rpc_addr);
~FileStorage();
tl::expected<void, ErrorCode> Init();
@ -104,10 +100,7 @@ class FileStorage {
void ClientBufferGCThreadFunc();
std::shared_ptr<Client> client_;
SsdMetric* ssd_metric_{nullptr};
std::string local_rpc_addr_;
// Pinned host memory pool for GPU D2H staging in OffloadObjects
std::unique_ptr<PinnedBufferPool> pinned_buffer_pool_;
std::shared_ptr<StorageBackendInterface> storage_backend_;
std::shared_ptr<ClientBufferAllocator> client_buffer_allocator_;
mutable Mutex client_buffer_mutex_;

View File

@ -1,76 +0,0 @@
#pragma once
#include "cuda_alike.h"
#if defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM)
#include <acl/acl_rt.h>
#endif
#include <cstddef>
#include <glog/logging.h>
namespace mooncake {
namespace gpu_staging {
// Detect whether ptr resides in accelerator device memory.
// If so, writes the device ID to *out_device_id for subsequent SetDevice.
inline bool IsDevicePointer(const void* ptr, int* out_device_id) {
#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA)
cudaPointerAttributes attr{};
if (cudaPointerGetAttributes(&attr, ptr) == cudaSuccess &&
attr.type == cudaMemoryTypeDevice) {
if (out_device_id) *out_device_id = attr.device;
return true;
}
#elif defined(USE_HIP)
hipPointerAttribute_t attr{};
if (hipPointerGetAttributes(&attr, ptr) == hipSuccess &&
attr.type == hipMemoryTypeDevice) {
if (out_device_id) *out_device_id = attr.device;
return true;
}
#elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM)
aclrtPtrAttributes attr{};
if (aclrtPointerGetAttributes(const_cast<void*>(ptr), &attr) ==
ACL_SUCCESS &&
attr.location.type == ACL_MEM_LOCATION_TYPE_DEVICE) {
if (out_device_id) *out_device_id = static_cast<int>(attr.location.id);
return true;
}
#endif
(void)ptr;
(void)out_device_id;
return false;
}
// Copy device memory to host. Caller must have called SetDevice first.
inline bool CopyDeviceToHost(void* dst, const void* src, size_t size) {
#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA)
return cudaMemcpy(dst, src, size, cudaMemcpyDeviceToHost) == cudaSuccess;
#elif defined(USE_HIP)
return hipMemcpy(dst, src, size, hipMemcpyDeviceToHost) == hipSuccess;
#elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM)
return aclrtMemcpy(dst, size, src, size, ACL_MEMCPY_DEVICE_TO_HOST) ==
ACL_SUCCESS;
#else
(void)dst;
(void)src;
(void)size;
return false;
#endif
}
// Bind the calling thread to the given device context.
inline void SetDevice(int device_id) {
if (device_id < 0) return;
#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA)
cudaSetDevice(device_id);
#elif defined(USE_HIP)
hipSetDevice(device_id);
#elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM)
aclrtSetDevice(device_id);
#endif
}
} // namespace gpu_staging
} // namespace mooncake

View File

@ -81,10 +81,6 @@ struct MasterConfig {
std::string cxl_path;
size_t cxl_size;
bool enable_cxl = false;
// Offload-on-evict: defer LOCAL_DISK offload to eviction time
bool offload_on_evict = false;
bool offload_force_evict = false;
};
class MasterServiceSupervisorConfig {
@ -144,8 +140,6 @@ class MasterServiceSupervisorConfig {
std::string cxl_path = DEFAULT_CXL_PATH;
size_t cxl_size = DEFAULT_CXL_SIZE;
bool enable_cxl = false;
bool offload_on_evict = false;
bool offload_force_evict = false;
MasterServiceSupervisorConfig() = default;
// From MasterConfig
@ -161,8 +155,6 @@ class MasterServiceSupervisorConfig {
eviction_high_watermark_ratio = config.eviction_high_watermark_ratio;
client_live_ttl_sec = config.client_live_ttl_sec;
enable_offload = config.enable_offload;
offload_on_evict = config.offload_on_evict;
offload_force_evict = config.offload_force_evict;
rpc_port = static_cast<int>(config.rpc_port);
rpc_thread_num = static_cast<size_t>(config.rpc_thread_num);
@ -275,8 +267,6 @@ class WrappedMasterServiceConfig {
int64_t client_live_ttl_sec = DEFAULT_CLIENT_LIVE_TTL_SEC;
bool enable_ha = false;
bool enable_offload = false;
bool offload_on_evict = false;
bool offload_force_evict = false;
std::string ha_backend_type = "etcd";
std::string ha_backend_connstring;
std::string cluster_id = DEFAULT_CLUSTER_ID;
@ -332,8 +322,6 @@ class WrappedMasterServiceConfig {
client_live_ttl_sec = config.client_live_ttl_sec;
enable_ha = config.enable_ha;
enable_offload = config.enable_offload;
offload_on_evict = config.offload_on_evict;
offload_force_evict = config.offload_force_evict;
ha_backend_type = config.ha_backend_type;
ha_backend_connstring = config.ha_backend_connstring;
if (ha_backend_connstring.empty()) {
@ -412,8 +400,6 @@ class WrappedMasterServiceConfig {
enable_ha =
true; // This is used in HA mode, so enable_ha should be true
enable_offload = config.enable_offload;
offload_on_evict = config.offload_on_evict;
offload_force_evict = config.offload_force_evict;
ha_backend_type = config.ha_backend_type;
ha_backend_connstring = config.ha_backend_connstring;
if (ha_backend_connstring.empty()) {
@ -750,8 +736,6 @@ class MasterServiceConfig {
int64_t client_live_ttl_sec = DEFAULT_CLIENT_LIVE_TTL_SEC;
bool enable_ha = false;
bool enable_offload = false;
bool offload_on_evict = false;
bool offload_force_evict = false;
std::string ha_backend_type = "etcd";
std::string ha_backend_connstring;
std::string cluster_id = DEFAULT_CLUSTER_ID;
@ -803,8 +787,6 @@ class MasterServiceConfig {
client_live_ttl_sec = config.client_live_ttl_sec;
enable_ha = config.enable_ha;
enable_offload = config.enable_offload;
offload_on_evict = config.offload_on_evict;
offload_force_evict = config.offload_force_evict;
ha_backend_type = config.ha_backend_type;
ha_backend_connstring = config.ha_backend_connstring;
cluster_id = config.cluster_id;

View File

@ -562,14 +562,8 @@ class MasterService {
// fulfill evict ratio lowerbound.
void BatchEvict(double evict_ratio_target, double evict_ratio_lowerbound);
// Helper to get a snapshot of alive clients (under client_mutex_ shared
// lock)
std::unordered_set<UUID, boost::hash<UUID>> getAliveClientsSnapshot() const;
// Clear invalid handles in all shards
void ClearInvalidHandles();
void ClearInvalidHandles(
const std::unordered_set<UUID, boost::hash<UUID>>& alive_clients);
std::string FormatTimestamp(
const std::chrono::system_clock::time_point& tp);
@ -897,10 +891,7 @@ class MasterService {
}
// Helper to clean up stale handles pointing to unmounted segments
// or local_disk replicas whose owner client is no longer alive.
bool CleanupStaleHandles(
ObjectMetadata& metadata,
const std::unordered_set<UUID, boost::hash<UUID>>& alive_clients);
bool CleanupStaleHandles(ObjectMetadata& metadata);
// Helper: allocate replicas, create ObjectMetadata, insert into shard,
// and return descriptor list. Shared by PutStart and UpsertStart.
@ -971,21 +962,11 @@ class MasterService {
it_(shard_guard_->metadata.find(key)),
processing_it_(shard_guard_->processing_keys.find(key)),
replication_task_it_(shard_guard_->replication_tasks.find(key)) {
// Automatically clean up invalid handles (memory replicas only).
// Note: We only check memory replicas here to avoid lock order
// violation (client_mutex_ must be acquired before metadata shard).
// local_disk replicas are cleaned up by ClearInvalidHandles() in
// ClientMonitorFunc.
// Automatically clean up invalid handles
if (it_ != shard_guard_->metadata.end()) {
// Erase invalid memory replicas (those with unmounted
// segments). No client_mutex_ needed since we only check memory
// replicas.
it_->second.EraseReplicas([](const Replica& replica) {
return replica.has_invalid_mem_handle();
});
// If no valid replicas remain, delete the whole object.
if (!it_->second.IsValid()) {
if (service_->CleanupStaleHandles(it_->second)) {
this->Erase();
if (processing_it_ != shard_guard_->processing_keys.end()) {
this->EraseFromProcessing();
}
@ -1168,14 +1149,6 @@ class MasterService {
const bool enable_offload_;
// Offload-on-evict: defer disk offload to eviction time
// (config: offload_on_evict)
bool offload_on_evict_{false};
// Force-evict: allow evicting MEMORY replicas without disk offload when cap
// exceeded (config: offload_force_evict, only effective when
// offload_on_evict_=true)
bool offload_force_evict_{false};
const std::string ha_backend_type_;
const std::string ha_backend_connstring_;

View File

@ -1,138 +0,0 @@
#pragma once
#include <mutex>
#include <vector>
#include <cstdlib>
#include "cuda_alike.h"
// Ascend CANN is not covered by cuda_alike.h
#if defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM)
#include <acl/acl_rt.h>
#endif
namespace mooncake {
/**
* PinnedBufferPool: Thread-safe pool of reusable pinned host memory buffers.
*
* Platform pinned alloc APIs:
* CUDA / MUSA / MACA : cudaMallocHost (mapped via cuda_alike.h)
* HIP : hipHostMalloc (not mapped in hip.h, native API)
* Ascend : aclrtMallocHost
* Other : new char[] (pageable fallback)
*
* Pinned memory provides 10x~100x higher D2H bandwidth than pageable memory.
* Falls back to new char[] if pinned allocation fails.
*
* The pool enforces a maximum number of cached buffers (kDefaultMaxPoolSize).
* When the pool is full, Release() frees the buffer immediately instead of
* caching it, preventing unbounded pinned memory growth.
*/
class PinnedBufferPool {
public:
static constexpr size_t kDefaultMaxPoolSize = 32;
struct Buffer {
char* data = nullptr;
size_t capacity = 0;
bool is_pinned = false; // Selects correct free API in FreeBuffer
};
explicit PinnedBufferPool(size_t max_pool_size = kDefaultMaxPoolSize)
: max_pool_size_(max_pool_size) {}
~PinnedBufferPool() { Clear(); }
Buffer Acquire(size_t size) {
{
std::lock_guard<std::mutex> lk(mutex_);
for (size_t i = 0; i < pool_.size(); ++i) {
if (pool_[i].capacity >= size) {
Buffer buf = pool_[i];
// O(1) erase: swap with back then pop
pool_[i] = pool_.back();
pool_.pop_back();
return buf;
}
}
}
return AllocNew(size);
}
void Release(Buffer buf) {
std::lock_guard<std::mutex> lk(mutex_);
if (pool_.size() < max_pool_size_) {
pool_.push_back(buf);
} else {
// Pool full — free immediately to bound pinned memory usage
FreeBuffer(buf);
}
}
void Clear() {
std::lock_guard<std::mutex> lk(mutex_);
for (auto& buf : pool_) {
FreeBuffer(buf);
}
pool_.clear();
}
private:
static Buffer AllocNew(size_t size) {
Buffer buf;
buf.capacity = size;
buf.is_pinned = false;
#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA)
if (cudaMallocHost(reinterpret_cast<void**>(&buf.data), size) ==
cudaSuccess) {
buf.is_pinned = true;
} else {
buf.data = new char[size];
}
#elif defined(USE_HIP)
if (hipHostMalloc(reinterpret_cast<void**>(&buf.data), size, 0) ==
hipSuccess) {
buf.is_pinned = true;
} else {
buf.data = new char[size];
}
#elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM)
if (aclrtMallocHost(reinterpret_cast<void**>(&buf.data), size) ==
ACL_SUCCESS) {
buf.is_pinned = true;
} else {
buf.data = new char[size];
}
#else
buf.data = new char[size];
#endif
return buf;
}
static void FreeBuffer(Buffer& buf) {
if (!buf.data) return;
if (!buf.is_pinned) {
delete[] buf.data;
return;
}
#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA)
cudaFreeHost(buf.data);
#elif defined(USE_HIP)
hipHostFree(buf.data);
#elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM)
aclrtFreeHost(buf.data);
#else
delete[] buf.data;
#endif
}
const size_t max_pool_size_;
std::mutex mutex_;
std::vector<Buffer> pool_;
};
} // namespace mooncake

View File

@ -211,8 +211,8 @@ class PyClient {
const std::string &protocol, const std::string &rdma_devices,
const std::string &master_server_addr,
const std::shared_ptr<TransferEngine> &transfer_engine,
const std::string &ipc_socket_path, bool enable_ssd_offload = false,
const std::string &ssd_offload_path = "") = 0;
const std::string &ipc_socket_path,
bool enable_ssd_offload = false) = 0;
virtual int setup_dummy(size_t mem_pool_size, size_t local_buffer_size,
const std::string &server_address,
@ -330,10 +330,6 @@ class PyClient {
virtual std::vector<Replica::Descriptor> get_replica_desc(
const std::string &key) = 0;
virtual std::vector<std::string> batch_replica_clear(
const std::vector<std::string> &keys,
const std::string &segment_name = "") = 0;
virtual int tearDownAll() = 0;
virtual int health_check() = 0;

View File

@ -79,8 +79,7 @@ class RealClient : public PyClient {
const std::string &master_server_addr = "127.0.0.1:50051",
const std::shared_ptr<TransferEngine> &transfer_engine = nullptr,
const std::string &ipc_socket_path = "",
bool enable_ssd_offload = false,
const std::string &ssd_offload_path = "");
bool enable_ssd_offload = false);
int setup_dummy(size_t mem_pool_size, size_t local_buffer_size,
const std::string &server_address,
@ -484,8 +483,7 @@ class RealClient : public PyClient {
const std::string &master_server_addr = "127.0.0.1:50051",
const std::shared_ptr<TransferEngine> &transfer_engine = nullptr,
const std::string &ipc_socket_path = "", int local_rpc_port = 50052,
bool enable_ssd_offload = false, bool start_offload_rpc_server = false,
const std::string &ssd_offload_path = "");
bool enable_ssd_offload = false, bool start_offload_rpc_server = false);
// Overload that accepts a configuration dictionary
tl::expected<void, ErrorCode> setup_internal(const ConfigDict &config);
@ -639,10 +637,6 @@ class RealClient : public PyClient {
batch_get_replica_desc(const std::vector<std::string> &keys);
std::vector<Replica::Descriptor> get_replica_desc(const std::string &key);
std::vector<std::string> batch_replica_clear(
const std::vector<std::string> &keys,
const std::string &segment_name = "") override;
tl::expected<PingResponse, ErrorCode> ping(const UUID &client_id);
tl::expected<BatchGetOffloadObjectResponse, ErrorCode>

View File

@ -2,12 +2,9 @@
#include <glog/logging.h>
#include <boost/functional/hash.hpp>
#include <cstdint>
#include <memory>
#include <string>
#include <unordered_set>
#include <variant>
#include <vector>
#include <unordered_map>
@ -285,36 +282,6 @@ class Replica {
return false; // DiskReplicaData does not have handles
}
/**
* @brief Check if a local_disk replica's owner client is still alive.
* Used by CleanupStaleHandles to remove replicas belonging to expired
* clients. For non-local_disk replicas, always returns false.
* @param alive_clients Set of currently alive client IDs.
* @return true if this is a local_disk replica whose client is not alive.
*/
[[nodiscard]] bool has_stale_local_disk_client(
const std::unordered_set<UUID, boost::hash<UUID>>& alive_clients)
const {
auto client_id = get_local_disk_client_id();
if (client_id.has_value()) {
return alive_clients.find(client_id.value()) == alive_clients.end();
}
return false;
}
/**
* @brief Get the client_id for local_disk replicas.
* @return The client_id if this is a local_disk replica, std::nullopt
* otherwise.
*/
[[nodiscard]] std::optional<UUID> get_local_disk_client_id() const {
if (is_local_disk_replica()) {
const auto& disk_data = std::get<LocalDiskReplicaData>(data_);
return disk_data.client_id;
}
return std::nullopt;
}
[[nodiscard]] size_t get_memory_buffer_size() const {
if (is_memory_replica()) {
const auto& mem_data = std::get<MemoryReplicaData>(data_);

View File

@ -170,12 +170,6 @@ class ScopedSegmentAccess {
ErrorCode SetSegmentStatusByName(const std::string& segment_name,
SegmentStatus status);
/**
* @brief Remove the local disk segment entry for a client.
* Called when a client expires to clean up its local disk segment.
*/
void UnmountLocalDiskSegment(const UUID& client_id);
private:
SegmentManager* segment_manager_;
std::unique_lock<std::shared_mutex> lock_;

View File

@ -174,10 +174,7 @@ class TransferEngineOperationState : public OperationState {
public:
TransferEngineOperationState(TransferEngine& engine, BatchID batch_id,
size_t batch_size)
: engine_(engine),
batch_id_(batch_id),
batch_size_(batch_size),
start_ts_(getCurrentTimeInMilli()) {}
: engine_(engine), batch_id_(batch_id), batch_size_(batch_size) {}
~TransferEngineOperationState() { engine_.freeBatchID(batch_id_); }
@ -202,7 +199,6 @@ class TransferEngineOperationState : public OperationState {
TransferEngine& engine_;
BatchID batch_id_;
size_t batch_size_;
const int64_t start_ts_;
};
/**

View File

@ -202,62 +202,15 @@ if(STORE_USE_ETCD)
add_dependencies(mooncake_master build_etcd_wrapper)
endif()
target_compile_options(mooncake_master PRIVATE -Os)
target_link_options(mooncake_master PRIVATE -Os -s)
# Client server binary
add_executable(mooncake_client real_client_main.cpp)
# Client needs transfer_engine for data transfer operations
target_link_libraries(mooncake_client PRIVATE mooncake_store transfer_engine
asio_shared)
target_compile_options(mooncake_client PRIVATE -Os)
target_link_options(mooncake_client PRIVATE -Os -s)
# Optimize binary sizes only in Release mode
string(TOUPPER "${CMAKE_BUILD_TYPE}" CMAKE_BUILD_TYPE_UPPER)
if (CMAKE_BUILD_TYPE_UPPER STREQUAL "RELEASE")
target_compile_options(mooncake_master PRIVATE -Os)
target_link_options(mooncake_master PRIVATE -Os -s)
target_compile_options(mooncake_client PRIVATE -Os)
target_link_options(mooncake_client PRIVATE -Os -s)
endif()
# GPU runtime library for D2H staging in PutToLocalFile / OffloadObjects.
# transfer_engine is PRIVATE-linked, so its CUDA/HIP/Ascend dependencies
# are not propagated; we must detect and link them independently.
#
# Auto-detect each toolkit regardless of global USE_CUDA/USE_HIP flags,
# because USE_CUDA may be OFF even when GPU pointers are present
# (e.g. WITH_NVIDIA_PEERMEM=ON uses nvidia-peermem for RDMA without cudart).
# Each detected toolkit gets both link libraries AND compile definitions,
# so that gpu_staging_utils.h / pinned_buffer_pool.h pick the correct backend.
#
# NOTE: mooncake_store is a static library (.a). External consumers (Go CGo,
# Python pybind) that link it must also link the GPU runtime (e.g. -lcudart).
# Go's build.sh already does this; CI workflows must do the same.
find_package(CUDAToolkit QUIET)
if(CUDAToolkit_FOUND)
message(STATUS "mooncake_store: CUDAToolkit detected, enabling D2H staging")
target_compile_definitions(mooncake_store PRIVATE USE_CUDA)
target_compile_definitions(mooncake_client PRIVATE USE_CUDA)
target_include_directories(mooncake_store PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
target_include_directories(mooncake_client PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
target_link_libraries(mooncake_store PRIVATE CUDA::cudart)
target_link_libraries(mooncake_client PRIVATE CUDA::cudart)
endif()
if(NOT CUDAToolkit_FOUND)
find_package(hip QUIET)
if(hip_FOUND)
message(STATUS "mooncake_store: HIP detected, enabling D2H staging")
target_compile_definitions(mooncake_store PRIVATE USE_HIP)
target_compile_definitions(mooncake_client PRIVATE USE_HIP)
target_link_libraries(mooncake_store PRIVATE hip::host)
target_link_libraries(mooncake_client PRIVATE hip::host)
endif()
endif()
if(USE_ASCEND OR USE_ASCEND_DIRECT OR USE_UBSHMEM)
target_include_directories(mooncake_store PRIVATE $ENV{ASCEND_PATH}/include)
target_link_libraries(mooncake_store PRIVATE ascendcl)
target_include_directories(mooncake_client PRIVATE $ENV{ASCEND_PATH}/include)
target_link_libraries(mooncake_client PRIVATE ascendcl)
endif()
install(TARGETS mooncake_master mooncake_client DESTINATION bin)

View File

@ -28,27 +28,6 @@ bool parseMetricsEnabled() {
value == "on" || value == "enable");
}
bool parseBoolEnv(const char* env_name, bool default_value) {
const char* env_value = std::getenv(env_name);
if (!env_value) {
return default_value;
}
std::string value = toLower(env_value);
if (value == "1" || value == "true" || value == "yes" || value == "on" ||
value == "enable") {
return true;
}
if (value == "0" || value == "false" || value == "no" || value == "off" ||
value == "disable") {
return false;
}
LOG(WARNING) << "Failed to parse " << env_name << ": " << env_value
<< ", fallback to default=" << default_value;
return default_value;
}
uint64_t parseMetricsInterval() {
const char* interval_env = std::getenv("MC_STORE_CLIENT_METRIC_INTERVAL");
if (!interval_env) {
@ -76,21 +55,11 @@ uint64_t parseMetricsInterval() {
} // anonymous namespace
ClientMetric::ClientMetric(uint64_t interval_seconds,
const std::map<std::string, std::string>& labels,
bool bandwidth_reporting_enabled,
bool master_rpc_metrics_enabled)
const std::map<std::string, std::string>& labels)
: transfer_metric(labels),
master_client_metric(labels),
transfer_operation_metric(labels),
ssd_metric(labels),
should_stop_metrics_thread_(false),
metrics_interval_seconds_(interval_seconds),
bandwidth_reporting_enabled_(bandwidth_reporting_enabled),
master_rpc_metrics_enabled_(master_rpc_metrics_enabled) {
last_report_snapshot_ = TransferSnapshot{
static_cast<uint64_t>(transfer_metric.total_read_bytes.value()),
static_cast<uint64_t>(transfer_metric.total_write_bytes.value()),
std::chrono::steady_clock::now()};
metrics_interval_seconds_(interval_seconds) {
if (metrics_interval_seconds_ > 0) {
StartMetricsReportingThread();
}
@ -99,8 +68,7 @@ ClientMetric::ClientMetric(uint64_t interval_seconds,
ClientMetric::~ClientMetric() { StopMetricsReportingThread(); }
std::unique_ptr<ClientMetric> ClientMetric::Create(
const std::map<std::string, std::string>& labels,
bool master_rpc_metrics_enabled) {
const std::map<std::string, std::string>& labels) {
if (!parseMetricsEnabled()) {
LOG(INFO) << "Client metrics disabled (set MC_STORE_CLIENT_METRIC=0 to "
"disable)";
@ -108,80 +76,23 @@ std::unique_ptr<ClientMetric> ClientMetric::Create(
}
uint64_t interval = parseMetricsInterval();
bool bandwidth_reporting_enabled =
parseBoolEnv("MC_STORE_CLIENT_METRIC_BANDWIDTH", true);
LOG(INFO) << "Client metrics enabled (default enabled)";
LOG(INFO) << "Client bandwidth summary "
<< (bandwidth_reporting_enabled ? "enabled" : "disabled")
<< " via MC_STORE_CLIENT_METRIC_BANDWIDTH";
return std::make_unique<ClientMetric>(interval, labels,
bandwidth_reporting_enabled,
master_rpc_metrics_enabled);
return std::make_unique<ClientMetric>(interval, labels);
}
void ClientMetric::serialize(std::string& str) {
transfer_metric.serialize(str);
if (master_rpc_metrics_enabled_) {
master_client_metric.serialize(str);
}
transfer_operation_metric.serialize(str);
ssd_metric.serialize(str);
master_client_metric.serialize(str);
}
std::string ClientMetric::summary_metrics() {
std::stringstream ss;
ss << "Client Metrics Summary\n";
ss << transfer_metric.summary_metrics(bandwidth_reporting_enabled_);
ss << transfer_metric.summary_metrics();
ss << "\n";
if (master_rpc_metrics_enabled_) {
ss << master_client_metric.summary_metrics();
ss << "\n";
}
ss << transfer_operation_metric.summary_metrics();
ss << "\n";
ss << ssd_metric.summary_metrics();
return ss.str();
}
std::string ClientMetric::BuildBandwidthReport() {
if (!bandwidth_reporting_enabled_) {
return "";
}
const auto now = std::chrono::steady_clock::now();
const uint64_t read_bytes = transfer_metric.total_read_bytes.value();
const uint64_t write_bytes = transfer_metric.total_write_bytes.value();
std::lock_guard<std::mutex> lock(snapshot_mutex_);
if (!last_report_snapshot_.has_value()) {
last_report_snapshot_ = TransferSnapshot{read_bytes, write_bytes, now};
return "";
}
const auto previous = *last_report_snapshot_;
last_report_snapshot_ = TransferSnapshot{read_bytes, write_bytes, now};
const double elapsed_seconds = std::max(
std::chrono::duration<double>(now - previous.timestamp).count(), 1e-9);
const uint64_t read_delta = read_bytes >= previous.read_bytes
? read_bytes - previous.read_bytes
: 0;
const uint64_t write_delta = write_bytes >= previous.write_bytes
? write_bytes - previous.write_bytes
: 0;
std::stringstream ss;
ss << "=== Interval Throughput Summary ===\n";
ss << "Read Throughput: "
<< format_metric_rate(read_delta / elapsed_seconds, "B/s") << " ("
<< byte_size_to_string(read_delta) << " over " << std::fixed
<< std::setprecision(2) << elapsed_seconds << "s)\n";
ss << "Write Throughput: "
<< format_metric_rate(write_delta / elapsed_seconds, "B/s") << " ("
<< byte_size_to_string(write_delta) << " over " << std::fixed
<< std::setprecision(2) << elapsed_seconds << "s)";
ss << master_client_metric.summary_metrics();
return ss.str();
}
@ -209,12 +120,7 @@ void ClientMetric::StartMetricsReportingThread() {
// Print metrics summary
std::string summary = summary_metrics();
std::string bandwidth_report = BuildBandwidthReport();
std::string report = "Client Metrics Report:\n" + summary;
if (!bandwidth_report.empty()) {
report += "\n" + bandwidth_report;
}
LOG(INFO) << report;
LOG(INFO) << "Client Metrics Report:\n" << summary;
}
LOG(INFO) << "Client metrics reporting thread stopped";
});
@ -230,4 +136,4 @@ void ClientMetric::StopMetricsReportingThread() {
}
}
} // namespace mooncake
} // namespace mooncake

View File

@ -27,14 +27,9 @@
#include "utils.h"
#include "rpc_types.h"
#include "local_hot_cache.h"
#include "gpu_staging_utils.h"
namespace mooncake {
using gpu_staging::CopyDeviceToHost;
using gpu_staging::IsDevicePointer;
using gpu_staging::SetDevice;
[[nodiscard]] size_t CalculateSliceSize(const std::vector<Slice>& slices) {
size_t slice_size = 0;
for (const auto& slice : slices) {
@ -62,7 +57,6 @@ Client::Client(const std::string& local_hostname,
local_hostname_(local_hostname),
metadata_connstring_(metadata_connstring),
protocol_(protocol),
pinned_buffer_pool_(std::make_unique<PinnedBufferPool>()),
write_thread_pool_(2),
task_thread_pool_(4) {
LOG(INFO) << "client_id=" << client_id_;
@ -2501,34 +2495,19 @@ void Client::PutToLocalFile(const std::string& key,
}
std::string path = disk_descriptor.file_path;
// Currently, persistence is achieved through asynchronous writes, but
// before asynchronous writing in 3FS, significant performance degradation
// may occur due to data copying. Profiling reveals that the number of page
// faults triggered in this scenario is nearly double the normal count.
// Future plans include introducing a reuse buffer list to address this
// performance degradation issue.
// Synchronous D2H staging + copy into std::string.
// Done on the calling thread to guarantee GPU buffers are still valid
// (BatchPut has not yet returned to Python, so blocks are not reused).
std::string value;
value.reserve(total_size);
for (const auto& slice : slices) {
int device_id = -1;
if (IsDevicePointer(slice.ptr, &device_id)) {
SetDevice(device_id);
auto buf = pinned_buffer_pool_->Acquire(slice.size);
if (!CopyDeviceToHost(buf.data, slice.ptr, slice.size)) {
LOG(ERROR) << "D2H copy failed for key: " << key
<< ", triggering PutRevoke for disk replica";
pinned_buffer_pool_->Release(buf);
// Must revoke to avoid phantom replica in master
master_client_.PutRevoke(key, ReplicaType::DISK);
return;
}
value.append(buf.data, slice.size);
pinned_buffer_pool_->Release(buf);
} else {
value.append(static_cast<char*>(slice.ptr), slice.size);
}
value.append(static_cast<char*>(slice.ptr), slice.size);
}
// Async StoreObject + PutEnd (unchanged from original)
write_thread_pool_.enqueue([this, backend = storage_backend_, key,
value = std::move(value), path] {
// Store the object

View File

@ -7,7 +7,6 @@
#include <sys/stat.h> // For S_IRUSR, S_IWUSR
#include <fcntl.h> // For O_CREAT, O_RDWR
#include <unistd.h> // For ftruncate, close, shm_unlink
#include <chrono>
#include <cstdlib>
#include "real_client.h"
@ -24,75 +23,6 @@
#endif
namespace {
size_t sum_value_sizes(const std::vector<std::span<const char>>& values) {
size_t total = 0;
for (const auto& value : values) {
total += value.size_bytes();
}
return total;
}
size_t sum_sizes(const std::vector<size_t>& sizes) {
size_t total = 0;
for (size_t size : sizes) {
total += size;
}
return total;
}
size_t sum_successful_sizes(const std::vector<int>& results,
const std::vector<size_t>& sizes) {
size_t total = 0;
for (size_t i = 0; i < results.size() && i < sizes.size(); ++i) {
if (results[i] == 0) {
total += sizes[i];
}
}
return total;
}
size_t sum_successful_nested_sizes(
const std::vector<int>& results,
const std::vector<std::vector<size_t>>& nested_sizes) {
size_t total = 0;
for (size_t i = 0; i < results.size() && i < nested_sizes.size(); ++i) {
if (results[i] == 0) {
total += sum_sizes(nested_sizes[i]);
}
}
return total;
}
size_t sum_positive_results(const std::vector<int64_t>& results) {
size_t total = 0;
for (int64_t result : results) {
if (result > 0) {
total += static_cast<size_t>(result);
}
}
return total;
}
size_t sum_positive_results(const std::vector<int>& results) {
size_t total = 0;
for (int result : results) {
if (result > 0) {
total += static_cast<size_t>(result);
}
}
return total;
}
size_t sum_positive_ranges(
const std::vector<std::vector<std::vector<int64_t>>>& results) {
size_t total = 0;
for (const auto& key_rows : results) {
for (const auto& row : key_rows) {
total += sum_positive_results(row);
}
}
return total;
}
std::vector<uint64_t> void_ptrs_to_u64(const std::vector<void*>& ptrs) {
std::vector<uint64_t> out;
@ -199,10 +129,7 @@ std::vector<tl::expected<ResultType, ErrorCode>> DummyClient::invoke_batch_rpc(
}());
}
DummyClient::DummyClient()
: client_id_(generate_uuid()),
metrics_(ClientMetric::Create(merge_labels({{"client_mode", "dummy"}}),
false)) {
DummyClient::DummyClient() : client_id_(generate_uuid()) {
// Initialize logging severity (leave as before)
mooncake::init_ylt_log_level();
// Initialize client pools
@ -214,30 +141,6 @@ DummyClient::DummyClient()
DummyClient::~DummyClient() { tearDownAll(); }
void DummyClient::ObserveTransferMetric(TransferOperationKind kind,
const char* op_name, size_t bytes,
uint64_t latency_us, bool batch) {
if (!metrics_) {
return;
}
metrics_->ObserveTransferOperation(kind, op_name, bytes, latency_us);
if (kind == TransferOperationKind::kRead) {
metrics_->transfer_metric.total_read_bytes.inc(bytes);
if (batch) {
metrics_->transfer_metric.batch_get_latency_us.observe(latency_us);
} else {
metrics_->transfer_metric.get_latency_us.observe(latency_us);
}
} else {
metrics_->transfer_metric.total_write_bytes.inc(bytes);
if (batch) {
metrics_->transfer_metric.batch_put_latency_us.observe(latency_us);
} else {
metrics_->transfer_metric.put_latency_us.observe(latency_us);
}
}
}
ErrorCode DummyClient::connect(const std::string& server_address) {
ScopedVLogTimer timer(1, "DummyClient::Connect");
timer.LogRequest("real_client_addr=", server_address);
@ -621,40 +524,35 @@ uint64_t DummyClient::alloc_from_mem_pool(size_t size) {
int DummyClient::put(const std::string& key, std::span<const char> value,
const ReplicateConfig& config) {
return invoke_observed_void_rpc<&RealClient::put_dummy_helper>(
TransferOperationKind::kWrite, "put", value.size_bytes(), false, key,
value, config, client_id_);
return to_py_ret(invoke_rpc<&RealClient::put_dummy_helper, void>(
key, value, config, client_id_));
}
int DummyClient::put_batch(const std::vector<std::string>& keys,
const std::vector<std::span<const char>>& values,
const ReplicateConfig& config) {
return invoke_observed_void_rpc<&RealClient::put_batch_dummy_helper>(
TransferOperationKind::kWrite, "put_batch", sum_value_sizes(values),
true, keys, values, config, client_id_);
return to_py_ret(invoke_rpc<&RealClient::put_batch_dummy_helper, void>(
keys, values, config, client_id_));
}
int DummyClient::put_parts(const std::string& key,
std::vector<std::span<const char>> values,
const ReplicateConfig& config) {
return invoke_observed_void_rpc<&RealClient::put_parts_dummy_helper>(
TransferOperationKind::kWrite, "put_parts", sum_value_sizes(values),
false, key, values, config, client_id_);
return to_py_ret(invoke_rpc<&RealClient::put_parts_dummy_helper, void>(
key, values, config, client_id_));
}
int DummyClient::upsert(const std::string& key, std::span<const char> value,
const ReplicateConfig& config) {
return invoke_observed_void_rpc<&RealClient::upsert_dummy_helper>(
TransferOperationKind::kWrite, "upsert", value.size_bytes(), false, key,
value, config, client_id_);
return to_py_ret(invoke_rpc<&RealClient::upsert_dummy_helper, void>(
key, value, config, client_id_));
}
int DummyClient::upsert_from(const std::string& key, void* buffer, size_t size,
const ReplicateConfig& config) {
uint64_t dummy_addr = reinterpret_cast<uint64_t>(buffer);
return invoke_observed_void_rpc<&RealClient::upsert_from_dummy_helper>(
TransferOperationKind::kWrite, "upsert_from", size, false, key,
dummy_addr, size, config, client_id_);
return to_py_ret(invoke_rpc<&RealClient::upsert_from_dummy_helper, void>(
key, dummy_addr, size, config, client_id_));
}
std::vector<int> DummyClient::batch_upsert_from(
@ -664,7 +562,6 @@ std::vector<int> DummyClient::batch_upsert_from(
for (auto ptr : buffer_ptrs) {
buffers.push_back(reinterpret_cast<uint64_t>(ptr));
}
const auto start_time = std::chrono::steady_clock::now();
auto internal_results =
invoke_batch_rpc<&RealClient::batch_upsert_from_dummy_helper, void>(
keys.size(), keys, buffers, sizes, config, client_id_);
@ -673,29 +570,21 @@ std::vector<int> DummyClient::batch_upsert_from(
for (const auto& result : internal_results) {
results.push_back(to_py_ret(result));
}
const size_t successful_bytes = sum_successful_sizes(results, sizes);
if (successful_bytes > 0) {
ObserveTransferMetric(TransferOperationKind::kWrite,
"batch_upsert_from", successful_bytes,
elapsed_us_since(start_time), true);
}
return results;
}
int DummyClient::upsert_parts(const std::string& key,
std::vector<std::span<const char>> values,
const ReplicateConfig& config) {
return invoke_observed_void_rpc<&RealClient::upsert_parts_dummy_helper>(
TransferOperationKind::kWrite, "upsert_parts", sum_value_sizes(values),
false, key, values, config, client_id_);
return to_py_ret(invoke_rpc<&RealClient::upsert_parts_dummy_helper, void>(
key, values, config, client_id_));
}
int DummyClient::upsert_batch(const std::vector<std::string>& keys,
const std::vector<std::span<const char>>& values,
const ReplicateConfig& config) {
return invoke_observed_void_rpc<&RealClient::upsert_batch_dummy_helper>(
TransferOperationKind::kWrite, "upsert_batch", sum_value_sizes(values),
true, keys, values, config, client_id_);
return to_py_ret(invoke_rpc<&RealClient::upsert_batch_dummy_helper, void>(
keys, values, config, client_id_));
}
int DummyClient::remove(const std::string& key, bool force) {
@ -761,7 +650,6 @@ int64_t DummyClient::getSize(const std::string& key) {
}
std::shared_ptr<BufferHandle> DummyClient::get_buffer(const std::string& key) {
const auto start_time = std::chrono::steady_clock::now();
// Try hot cache path if shm is mapped
if (hot_cache_base_) {
auto result = invoke_rpc<&RealClient::acquire_hot_cache,
@ -780,8 +668,6 @@ std::shared_ptr<BufferHandle> DummyClient::get_buffer(const std::string& key) {
(void)invoke_rpc<&RealClient::release_hot_cache, void>(
key_copy);
};
ObserveTransferMetric(TransferOperationKind::kRead, "get_buffer",
size, elapsed_us_since(start_time), false);
return std::make_shared<BufferHandle>(local_ptr, size,
std::move(release));
}
@ -800,14 +686,11 @@ std::shared_ptr<BufferHandle> DummyClient::get_buffer(const std::string& key) {
(void)invoke_rpc<&RealClient::release_buffer_dummy, void>(dummy_addr,
client_id_);
};
ObserveTransferMetric(TransferOperationKind::kRead, "get_buffer", size,
elapsed_us_since(start_time), false);
return std::make_shared<BufferHandle>(local_ptr, size, std::move(release));
}
std::vector<std::shared_ptr<BufferHandle>> DummyClient::batch_get_buffer(
const std::vector<std::string>& keys) {
const auto start_time = std::chrono::steady_clock::now();
std::vector<std::shared_ptr<BufferHandle>> results(keys.size(), nullptr);
if (keys.empty()) return results;
@ -868,37 +751,19 @@ std::vector<std::shared_ptr<BufferHandle>> DummyClient::batch_get_buffer(
std::make_shared<BufferHandle>(ptr, size, std::move(release));
}
size_t total_bytes = 0;
for (const auto& result : results) {
if (result != nullptr) {
total_bytes += result->size();
}
}
if (total_bytes > 0) {
ObserveTransferMetric(TransferOperationKind::kRead, "batch_get_buffer",
total_bytes, elapsed_us_since(start_time), true);
}
return results;
}
int64_t DummyClient::get_into(const std::string& key, void* buffer,
size_t size) {
uint64_t buf_addr = reinterpret_cast<uint64_t>(buffer);
const auto start_time = std::chrono::steady_clock::now();
auto result = invoke_rpc<&RealClient::get_into_range_shm_helper,
tl::expected<int64_t, ErrorCode>>(
key, buf_addr, 0, 0, size, client_id_);
if (!result) {
return static_cast<int64_t>(toInt(result.error()));
}
const int64_t bytes_read = to_py_ret(*result);
if (bytes_read >= 0) {
ObserveTransferMetric(TransferOperationKind::kRead, "get_into",
static_cast<size_t>(bytes_read),
elapsed_us_since(start_time), false);
}
return bytes_read;
return to_py_ret(*result);
}
std::vector<std::vector<std::vector<int64_t>>> DummyClient::get_into_ranges(
@ -908,7 +773,6 @@ std::vector<std::vector<std::vector<int64_t>>> DummyClient::get_into_ranges(
const std::vector<std::vector<std::vector<size_t>>>& all_src_offsets,
const std::vector<std::vector<std::vector<size_t>>>& all_sizes) {
std::vector<uint64_t> dummy_buffers = void_ptrs_to_u64(buffers);
const auto start_time = std::chrono::steady_clock::now();
auto internal_results =
invoke_rpc<&RealClient::get_into_ranges_shm_helper,
std::vector<std::vector<
@ -922,13 +786,8 @@ std::vector<std::vector<std::vector<int64_t>>> DummyClient::get_into_ranges(
all_dst_offsets,
internal_results.error());
}
auto results = convert_ranged_read_results(internal_results.value());
const size_t total_bytes = sum_positive_ranges(results);
if (total_bytes > 0) {
ObserveTransferMetric(TransferOperationKind::kRead, "get_into_ranges",
total_bytes, elapsed_us_since(start_time), true);
}
return results;
return convert_ranged_read_results(internal_results.value());
}
std::string DummyClient::get_hostname() const {
@ -940,7 +799,6 @@ std::vector<int> DummyClient::batch_put_from(
const std::vector<std::string>& keys, const std::vector<void*>& buffer_ptrs,
const std::vector<size_t>& sizes, const ReplicateConfig& config) {
std::vector<uint64_t> buffers = void_ptrs_to_u64(buffer_ptrs);
const auto start_time = std::chrono::steady_clock::now();
auto internal_results =
invoke_batch_rpc<&RealClient::batch_put_from_dummy_helper, void>(
keys.size(), keys, buffers, sizes, config, device_id_, client_id_);
@ -951,13 +809,6 @@ std::vector<int> DummyClient::batch_put_from(
results.push_back(to_py_ret(result));
}
const size_t successful_bytes = sum_successful_sizes(results, sizes);
if (successful_bytes > 0) {
ObserveTransferMetric(TransferOperationKind::kWrite, "batch_put_from",
successful_bytes, elapsed_us_since(start_time),
true);
}
return results;
}
@ -971,7 +822,6 @@ std::vector<int64_t> DummyClient::batch_get_into(
const std::vector<std::string>& keys, const std::vector<void*>& buffer_ptrs,
const std::vector<size_t>& sizes) {
std::vector<uint64_t> buffers = void_ptrs_to_u64(buffer_ptrs);
const auto start_time = std::chrono::steady_clock::now();
auto internal_results =
invoke_batch_rpc<&RealClient::batch_get_into_dummy_helper, int64_t>(
keys.size(), keys, buffers, sizes, device_id_, client_id_);
@ -982,12 +832,6 @@ std::vector<int64_t> DummyClient::batch_get_into(
results.push_back(to_py_ret(result));
}
const size_t total_bytes = sum_positive_results(results);
if (total_bytes > 0) {
ObserveTransferMetric(TransferOperationKind::kRead, "batch_get_into",
total_bytes, elapsed_us_since(start_time), true);
}
return results;
}
@ -1006,7 +850,6 @@ std::vector<int> DummyClient::batch_put_from_multi_buffers(
const ReplicateConfig& config) {
std::vector<std::vector<uint64_t>> dummy_nested =
void_ptr_rows_to_u64_nested(all_buffer_ptrs);
const auto start_time = std::chrono::steady_clock::now();
auto internal_results =
invoke_batch_rpc<&RealClient::batch_put_from_multi_buffers_dummy_helper,
void>(keys.size(), keys, dummy_nested, all_sizes,
@ -1016,13 +859,6 @@ std::vector<int> DummyClient::batch_put_from_multi_buffers(
for (const auto& result : internal_results) {
results.push_back(to_py_ret(result));
}
const size_t successful_bytes =
sum_successful_nested_sizes(results, all_sizes);
if (successful_bytes > 0) {
ObserveTransferMetric(TransferOperationKind::kWrite,
"batch_put_from_multi_buffers", successful_bytes,
elapsed_us_since(start_time), true);
}
return results;
}
@ -1033,7 +869,6 @@ std::vector<int> DummyClient::batch_get_into_multi_buffers(
bool prefer_alloc_in_same_node) {
std::vector<std::vector<uint64_t>> dummy_nested =
void_ptr_rows_to_u64_nested(all_buffer_ptrs);
const auto start_time = std::chrono::steady_clock::now();
auto internal_results =
invoke_batch_rpc<&RealClient::batch_get_into_multi_buffers_dummy_helper,
int64_t>(keys.size(), keys, dummy_nested, all_sizes,
@ -1044,12 +879,6 @@ std::vector<int> DummyClient::batch_get_into_multi_buffers(
for (const auto& result : internal_results) {
results.push_back(to_py_ret(result));
}
const size_t total_bytes = sum_positive_results(results);
if (total_bytes > 0) {
ObserveTransferMetric(TransferOperationKind::kRead,
"batch_get_into_multi_buffers", total_bytes,
elapsed_us_since(start_time), true);
}
return results;
}

View File

@ -5,19 +5,13 @@
#include "aligned_client_buffer.hpp"
#include "storage_backend.h"
#include "client_metric.h"
#include "utils.h"
#include "gpu_staging_utils.h"
#ifdef USE_URING
#include "file_interface.h"
#endif
namespace mooncake {
using gpu_staging::CopyDeviceToHost;
using gpu_staging::IsDevicePointer;
using gpu_staging::SetDevice;
FileStorageConfig FileStorageConfig::FromEnvironment() {
FileStorageConfig config;
@ -42,10 +36,9 @@ FileStorageConfig FileStorageConfig::FromEnvironment() {
config.local_buffer_size = GetEnvOr<int64_t>(
"MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES", config.local_buffer_size);
config.scanmeta_iterator_keys_limit = GetEnvOr<int64_t>(
"MOONCAKE_OFFLOAD_SCANMETA_ITERATOR_KEYS_LIMIT",
config.scanmeta_iterator_keys_limit =
GetEnvOr<int64_t>("MOONCAKE_SCANMETA_ITERATOR_KEYS_LIMIT",
config.scanmeta_iterator_keys_limit));
config.scanmeta_iterator_keys_limit);
config.total_keys_limit = GetEnvOr<int64_t>(
"MOONCAKE_OFFLOAD_TOTAL_KEYS_LIMIT", config.total_keys_limit);
@ -64,9 +57,7 @@ FileStorageConfig FileStorageConfig::FromEnvironment() {
GetEnvOr<uint64_t>("MOONCAKE_OFFLOAD_CLIENT_BUFFER_GC_TTL_MS",
config.client_buffer_gc_ttl_ms);
auto use_uring_str =
GetEnvStringOr("MOONCAKE_OFFLOAD_USE_URING",
GetEnvStringOr("MOONCAKE_USE_URING", "false"));
auto use_uring_str = GetEnvStringOr("MOONCAKE_USE_URING", "false");
config.use_uring = (use_uring_str == "true" || use_uring_str == "1");
return config;
@ -154,13 +145,10 @@ bool FileStorageConfig::Validate() const {
FileStorage::FileStorage(const FileStorageConfig& config,
std::shared_ptr<Client> client,
const std::string& local_rpc_addr,
SsdMetric* ssd_metric)
const std::string& local_rpc_addr)
: config_(config),
client_(client),
ssd_metric_(ssd_metric),
local_rpc_addr_(local_rpc_addr),
pinned_buffer_pool_(std::make_unique<PinnedBufferPool>()),
client_buffer_allocator_(
AlignedClientBufferAllocator::create(config.local_buffer_size, "")) {
if (!config.Validate()) {
@ -381,70 +369,8 @@ tl::expected<void, ErrorCode> FileStorage::OffloadObjects(
}
};
// D2H staging: replace device slices with host memory slices
// so that storage_backend (ConcatSlicesToString / BuildBucket /
// WriteBucket) always receives host pointers.
std::unordered_map<std::string, std::vector<Slice>> host_batch_object;
std::vector<PinnedBufferPool::Buffer> staging_bufs;
for (auto& [obj_key, slices] : batch_object) {
std::vector<Slice> host_slices;
bool obj_success = true;
for (const auto& slice : slices) {
int device_id = -1;
if (IsDevicePointer(slice.ptr, &device_id)) {
SetDevice(device_id);
auto buf = pinned_buffer_pool_->Acquire(slice.size);
if (!CopyDeviceToHost(buf.data, slice.ptr, slice.size)) {
LOG(ERROR) << "D2H staging failed for key: " << obj_key;
pinned_buffer_pool_->Release(buf);
obj_success = false;
break;
}
host_slices.emplace_back(Slice{buf.data, slice.size});
staging_bufs.push_back(buf);
} else {
host_slices.push_back(slice);
}
}
if (obj_success) {
host_batch_object[obj_key] = std::move(host_slices);
}
}
auto offload_start = std::chrono::steady_clock::now();
auto bucket_complete_handler =
[this, offload_start, complete_handler](
const std::vector<std::string>& keys,
std::vector<StorageObjectMetadata>& metadatas) -> ErrorCode {
auto res = complete_handler(keys, metadatas);
if (res == ErrorCode::OK && ssd_metric_) {
auto elapsed_us =
std::chrono::duration_cast<std::chrono::microseconds>(
std::chrono::steady_clock::now() - offload_start)
.count();
int64_t total_bytes = 0;
for (const auto& metadata : metadatas) {
total_bytes += metadata.data_size;
}
ssd_metric_->ssd_write_ops.inc(keys.size());
ssd_metric_->ssd_write_bytes.inc(total_bytes);
ssd_metric_->ssd_write_latency_us.observe(elapsed_us);
ssd_metric_->ssd_write_latency_summary.observe(elapsed_us);
ssd_metric_->ssd_total_ops.inc(keys.size());
ssd_metric_->ssd_total_bytes.inc(total_bytes);
ssd_metric_->ssd_total_latency_us.observe(elapsed_us);
ssd_metric_->ssd_total_latency_summary.observe(elapsed_us);
}
return res;
};
auto offload_res = storage_backend_->BatchOffload(
host_batch_object, bucket_complete_handler, eviction_handler);
// Release staging buffers back to pool (Buffer is POD, no destructor)
for (auto& buf : staging_bufs) {
pinned_buffer_pool_->Release(buf);
}
batch_object, complete_handler, eviction_handler);
if (!offload_res) {
LOG(ERROR) << "Failed to store objects with error: "
<< offload_res.error();
@ -519,19 +445,6 @@ tl::expected<void, ErrorCode> FileStorage::BatchLoad(
<< "us,with keys count: " << batch_object.size();
if (!result) {
LOG(ERROR) << "Batch load object failed,err_code = " << result.error();
} else if (ssd_metric_) {
int64_t total_bytes = 0;
for (const auto& [key, slice] : batch_object) {
total_bytes += slice.size;
}
ssd_metric_->ssd_read_ops.inc(batch_object.size());
ssd_metric_->ssd_read_bytes.inc(total_bytes);
ssd_metric_->ssd_read_latency_us.observe(elapsed_time);
ssd_metric_->ssd_read_latency_summary.observe(elapsed_time);
ssd_metric_->ssd_total_ops.inc(batch_object.size());
ssd_metric_->ssd_total_bytes.inc(total_bytes);
ssd_metric_->ssd_total_latency_us.observe(elapsed_time);
ssd_metric_->ssd_total_latency_summary.observe(elapsed_time);
}
return result;
}

View File

@ -109,10 +109,6 @@ DEFINE_validator(eviction_ratio, [](const char* flagname, double value) {
DEFINE_bool(enable_ha, false,
"Enable high availability, which depends on etcd");
DEFINE_bool(enable_offload, false, "Enable offload availability");
DEFINE_bool(offload_on_evict, false,
"Defer LOCAL_DISK offload to eviction time instead of PutEnd");
DEFINE_bool(offload_force_evict, false,
"Force-evict objects exceeding offload cap without disk offload");
DEFINE_string(ha_backend_type, "etcd",
"HA backend type, e.g. etcd | redis | k8s");
DEFINE_string(ha_backend_connstring, "",
@ -297,11 +293,6 @@ void InitMasterConf(const mooncake::DefaultConfig& default_config,
FLAGS_enable_ha);
default_config.GetBool("enable_offload", &master_config.enable_offload,
FLAGS_enable_offload);
default_config.GetBool("offload_on_evict", &master_config.offload_on_evict,
FLAGS_offload_on_evict);
default_config.GetBool("offload_force_evict",
&master_config.offload_force_evict,
FLAGS_offload_force_evict);
default_config.GetString("ha_backend_type", &master_config.ha_backend_type,
FLAGS_ha_backend_type);
default_config.GetString("ha_backend_connstring",
@ -876,8 +867,6 @@ int main(int argc, char* argv[]) {
<< master_config.eviction_high_watermark_ratio
<< ", enable_ha=" << master_config.enable_ha
<< ", enable_offload=" << master_config.enable_offload
<< ", offload_on_evict=" << master_config.offload_on_evict
<< ", offload_force_evict=" << master_config.offload_force_evict
<< ", ha_backend_type=" << master_config.ha_backend_type
<< ", ha_backend_connstring=" << ha_backend_connstring
<< ", etcd_endpoints=" << master_config.etcd_endpoints

View File

@ -48,13 +48,6 @@ namespace {
constexpr size_t kUnlimitedSnapshotList = 0;
// Per-cycle offload cap as a fraction of `offloading_queue_limit_`. Used only
// when offload-on-evict mode is active. Defers memory eviction for at most
// this fraction of the queue limit per BatchEvict cycle; beyond that, eviction
// falls back according to `offload_force_evict_`. A future change may expose
// this as a configurable parameter if workloads demand tuning.
constexpr double kOffloadCapRatio = 0.5;
enum class SnapshotCatalogBackendKind {
kEmbedded,
kRedis,
@ -166,19 +159,6 @@ MasterService::MasterService(const MasterServiceConfig& config)
"put_start_discard_timeout_sec");
}
// Offload-on-evict: defer LOCAL_DISK offload to eviction time
offload_on_evict_ = enable_offload_ && config.offload_on_evict;
if (offload_on_evict_) {
LOG(INFO) << "Offload-on-evict mode enabled: DRAM offload to "
"LOCAL_DISK will occur at eviction time instead of "
"PutEnd";
offload_force_evict_ = config.offload_force_evict;
if (offload_force_evict_) {
LOG(INFO) << "Force-evict enabled: objects exceeding offload "
"cap will be evicted without disk offload";
}
}
eviction_running_ = true;
eviction_thread_ = std::thread(&MasterService::EvictionThreadFunc, this);
VLOG(1) << "action=start_eviction_thread";
@ -371,23 +351,12 @@ auto MasterService::ReMountSegment(const std::vector<Segment>& segments,
return {};
}
std::unordered_set<UUID, boost::hash<UUID>>
MasterService::getAliveClientsSnapshot() const {
std::shared_lock<std::shared_mutex> lock(client_mutex_);
return ok_client_;
}
void MasterService::ClearInvalidHandles() {
ClearInvalidHandles(getAliveClientsSnapshot());
}
void MasterService::ClearInvalidHandles(
const std::unordered_set<UUID, boost::hash<UUID>>& alive_clients) {
for (size_t i = 0; i < kNumShards; i++) {
MetadataShardAccessorRW shard(this, i);
auto it = shard->metadata.begin();
while (it != shard->metadata.end()) {
if (CleanupStaleHandles(it->second, alive_clients)) {
if (CleanupStaleHandles(it->second)) {
// If the object is empty, we need to erase the iterator and
// also erase the key from processing_keys,
// replication_tasks, and offloading_tasks.
@ -873,15 +842,13 @@ auto MasterService::PutStart(const UUID& client_id, const std::string& key,
VLOG(1) << "key=" << key << ", value_length=" << slice_length
<< ", config=" << config << ", action=put_start_begin";
auto alive_clients = getAliveClientsSnapshot();
std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
// Lock the shard and check if object already exists
MetadataShardAccessorRW shard(this, getShardIndex(key));
const auto now = std::chrono::system_clock::now();
auto it = shard->metadata.find(key);
if (it != shard->metadata.end() &&
!CleanupStaleHandles(it->second, alive_clients)) {
if (it != shard->metadata.end() && !CleanupStaleHandles(it->second)) {
auto& metadata = it->second;
// If the object's PutStart expired and has not completed any
// replicas, we can discard it and allow the new PutStart to
@ -930,7 +897,7 @@ auto MasterService::PutEnd(const UUID& client_id, const std::string& key,
},
[](Replica& replica) { replica.mark_complete(); });
if (enable_offload_ && !offload_on_evict_) {
if (enable_offload_) {
auto& shard = accessor.GetShard();
metadata.VisitReplicas(
&Replica::fn_is_completed, [this, &key, &shard](Replica& replica) {
@ -1118,7 +1085,6 @@ auto MasterService::UpsertStart(const UUID& client_id, const std::string& key,
// during full metadata snapshots.
// shard lock (exclusive via MetadataShardAccessorRW): serializes all
// operations on keys that hash to the same shard.
auto alive_clients = getAliveClientsSnapshot();
std::shared_lock<std::shared_mutex> shared_lock(snapshot_mutex_);
MetadataShardAccessorRW shard(this, getShardIndex(key));
@ -1128,9 +1094,7 @@ auto MasterService::UpsertStart(const UUID& client_id, const std::string& key,
// --- Step 0: stale handle cleanup ---
// If all memory replicas point to unmounted segments (node crashed and
// restarted), the metadata is useless — erase it and treat as new key.
// Also clean up local_disk replicas whose owner client has expired.
if (it != shard->metadata.end() &&
CleanupStaleHandles(it->second, alive_clients)) {
if (it != shard->metadata.end() && CleanupStaleHandles(it->second)) {
shard->processing_keys.erase(key);
shard->metadata.erase(it);
it = shard->metadata.end();
@ -1967,8 +1931,6 @@ auto MasterService::BatchRemove(const std::vector<std::string>& keys,
std::shared_lock<std::shared_mutex> snapshot_lock(snapshot_mutex_);
auto alive_clients = getAliveClientsSnapshot();
// Process each shard once, acquiring lock per shard
for (auto& [shard_idx, key_group] : keys_by_shard) {
MetadataShardAccessorRW shard(this, shard_idx);
@ -1986,7 +1948,7 @@ auto MasterService::BatchRemove(const std::vector<std::string>& keys,
}
// Clean up stale replica handles (consistent with single Remove)
if (CleanupStaleHandles(it->second, alive_clients)) {
if (CleanupStaleHandles(it->second)) {
shard->processing_keys.erase(key);
shard->replication_tasks.erase(key);
shard->offloading_tasks.erase(key);
@ -2034,14 +1996,10 @@ auto MasterService::BatchRemove(const std::vector<std::string>& keys,
return results;
}
bool MasterService::CleanupStaleHandles(
ObjectMetadata& metadata,
const std::unordered_set<UUID, boost::hash<UUID>>& alive_clients) {
// Remove those with invalid allocators (memory replicas on unmounted
// segments) and local_disk replicas whose owner client is no longer alive.
metadata.EraseReplicas([&alive_clients](const Replica& replica) {
return replica.has_invalid_mem_handle() ||
replica.has_stale_local_disk_client(alive_clients);
bool MasterService::CleanupStaleHandles(ObjectMetadata& metadata) {
// Remove those with invalid allocators
metadata.EraseReplicas([](const Replica& replica) {
return replica.has_invalid_mem_handle();
});
// Return true if no valid replicas remain after cleanup
@ -3511,81 +3469,6 @@ void MasterService::BatchEvict(double evict_ratio_target,
});
};
// --- Offload-on-evict support ---
long offload_queued_this_cycle = 0;
long offload_deferred_count = 0;
long offload_cap_forced_count = 0; // #keys force-evicted due to cap
long offload_push_failed_forced = 0; // #keys force-evicted on push fail
const long offload_cap =
offload_on_evict_
? static_cast<long>(offloading_queue_limit_ * kOffloadCapRatio)
: 0;
auto has_local_disk_replica = [](const ObjectMetadata& metadata) {
return metadata.HasReplica(&Replica::fn_is_local_disk_replica);
};
// Returns freed bytes. Returns 0 if offload-queued and no additional
// replicas were evicted (all MEMORY replicas of the key are now pinned).
auto try_evict_or_offload =
[&, this](const std::string& key, ObjectMetadata& metadata,
MetadataShardAccessorRW& shard) -> uint64_t {
if (!offload_on_evict_) {
// Original behavior
return metadata.size * evict_replicas(metadata);
}
// LOCAL_DISK replica already exists — safe to delete MEMORY immediately
if (has_local_disk_replica(metadata)) {
return metadata.size * evict_replicas(metadata);
}
// Force-evict cap: if force_evict enabled and cap reached, force
// delete. Warning is aggregated at the end of the cycle to avoid log
// flooding.
if (offload_force_evict_ && offload_queued_this_cycle >= offload_cap) {
offload_cap_forced_count++;
return metadata.size * evict_replicas(metadata);
}
// Queue one MEMORY replica for offload; others will be evicted below.
bool queued = false;
metadata.VisitReplicas(
[](const Replica& r) {
return r.is_memory_replica() && r.is_completed() &&
r.get_refcnt() == 0;
},
[this, &key, &shard, &queued, &now](Replica& replica) {
if (queued) return; // only need to pin one replica for offload
auto result = PushOffloadingQueue(key, replica);
if (result) {
replica.inc_refcnt();
shard->offloading_tasks.emplace(
key, OffloadingTask{replica.id(), now});
queued = true;
}
});
if (queued) {
offload_queued_this_cycle++;
offload_deferred_count++;
// Any remaining MEMORY replicas with refcnt==0 are redundant copies
// (data survives via the pinned replica → disk). Evict them now to
// reclaim memory immediately rather than waiting another cycle.
return metadata.size * evict_replicas(metadata);
}
// PushOffloadingQueue failed. Default (data-preserving) behavior is to
// skip this cycle — the outer eviction loop will retry after the
// offload queue drains. Only force-evict when explicitly opted in, to
// prevent silent data loss when the queue is unavailable.
if (offload_force_evict_) {
offload_push_failed_forced++;
return metadata.size * evict_replicas(metadata);
}
return 0;
};
// Randomly select a starting shard to avoid imbalance eviction between
// shards. No need to use expensive random_device here.
size_t start_idx = rand() % kNumShards;
@ -3658,18 +3541,16 @@ void MasterService::BatchEvict(double evict_ratio_target,
continue;
}
if (it->second.lease_timeout <= target_timeout) {
// Evict this object (or defer for offload)
uint64_t freed =
try_evict_or_offload(it->first, it->second, shard);
total_freed_size += freed;
// Evict this object
total_freed_size +=
it->second.size *
evict_replicas(it->second); // Erase memory replicas
if (it->second.IsValid() == false) {
it = shard->metadata.erase(it);
} else {
++it;
}
if (freed > 0) {
shard_evicted_count++;
}
shard_evicted_count++;
} else {
// second pass candidates
no_pin_objects.push_back(it->second.lease_timeout);
@ -3717,22 +3598,20 @@ void MasterService::BatchEvict(double evict_ratio_target,
auto it = shard->metadata.begin();
while (it != shard->metadata.end() && target_evict_num > 0) {
if (!it->second.IsHardPinned() &&
it->second.IsLeaseExpired(now) &&
it->second.lease_timeout <= target_timeout &&
!it->second.IsSoftPinned(now) &&
can_evict_replicas(it->second)) {
// Evict this object (or defer for offload)
uint64_t freed =
try_evict_or_offload(it->first, it->second, shard);
total_freed_size += freed;
// Evict this object
total_freed_size +=
it->second.size *
evict_replicas(
it->second); // Erase memory replicas
if (it->second.IsValid() == false) {
it = shard->metadata.erase(it);
} else {
++it;
}
if (freed > 0) {
evicted_count++;
}
evicted_count++;
target_evict_num--;
} else {
++it;
@ -3772,18 +3651,16 @@ void MasterService::BatchEvict(double evict_ratio_target,
// and lease timeout less than or equal to target.
if (!it->second.IsSoftPinned(now) ||
it->second.lease_timeout <= soft_target_timeout) {
// Evict this object (or defer for offload)
uint64_t freed =
try_evict_or_offload(it->first, it->second, shard);
total_freed_size += freed;
total_freed_size +=
it->second.size *
evict_replicas(
it->second); // Erase memory replicas
if (it->second.IsValid() == false) {
it = shard->metadata.erase(it);
} else {
++it;
}
if (freed > 0) {
evicted_count++;
}
evicted_count++;
target_evict_num--;
} else {
++it;
@ -3804,11 +3681,7 @@ void MasterService::BatchEvict(double evict_ratio_target,
}
}
if (evicted_count > 0 || released_discarded_cnt > 0 ||
offload_deferred_count > 0) {
// Offload-deferred counts as partial success: work was done (objects
// queued for disk offload), so suppress re-triggering until the next
// watermark breach or explicit need_eviction_ signal.
if (evicted_count > 0 || released_discarded_cnt > 0) {
need_eviction_ = false;
MasterMetricManager::instance().inc_eviction_success(evicted_count,
total_freed_size);
@ -3820,27 +3693,7 @@ void MasterService::BatchEvict(double evict_ratio_target,
MasterMetricManager::instance().inc_eviction_fail();
}
VLOG(1) << "action=evict_objects" << ", evicted_count=" << evicted_count
<< ", offload_deferred=" << offload_deferred_count
<< ", offload_cap_forced=" << offload_cap_forced_count
<< ", offload_push_failed_forced=" << offload_push_failed_forced
<< ", total_freed_size=" << total_freed_size;
if (offload_on_evict_ && evicted_count == 0 && offload_deferred_count > 0) {
LOG(WARNING) << "[EVICT] No memory freed this cycle; "
<< offload_deferred_count
<< " objects deferred for disk offload. "
"Consider lowering eviction_high_watermark_ratio.";
}
if (offload_cap_forced_count > 0) {
LOG(WARNING) << "[EVICT] Offload cap (" << offload_cap
<< ") reached; force-evicted " << offload_cap_forced_count
<< " object(s) without disk offload this cycle.";
}
if (offload_push_failed_forced > 0) {
LOG(WARNING) << "[EVICT] PushOffloadingQueue failed for "
<< offload_push_failed_forced
<< " object(s); force-evicted without disk offload "
"(offload_force_evict=true).";
}
}
void MasterService::ClientMonitorFunc() {
@ -3917,15 +3770,9 @@ void MasterService::ClientMonitorFunc() {
} // Release the mutex before long-running ClearInvalidHandles and
// avoid deadlocks
// Always clean up invalid handles when there are expired clients,
// even if no memory segments were unmounted. This is necessary
// to clean up local_disk replicas whose owner client has expired.
ClearInvalidHandles();
if (!unmount_segments.empty()) {
ClearInvalidHandles();
// Commit unmount of memory segments and clean up local_disk
// segments for expired clients. Both require the exclusive
// segment lock.
{
ScopedSegmentAccess segment_access =
segment_manager_.getSegmentAccess();
for (size_t i = 0; i < unmount_segments.size(); i++) {
@ -3935,9 +3782,6 @@ void MasterService::ClientMonitorFunc() {
<< ", segment_name=" << segment_names[i]
<< ", action=unmount_expired_segment";
}
for (auto& client_id : expired_clients) {
segment_access.UnmountLocalDiskSegment(client_id);
}
}
}

View File

@ -62,87 +62,6 @@ struct PreparedRangedReadRequest {
bool has_any_valid_fragment = false;
};
size_t sum_value_sizes(const std::vector<std::span<const char>> &values) {
size_t total = 0;
for (const auto &value : values) {
total += value.size_bytes();
}
return total;
}
size_t sum_sizes(const std::vector<size_t> &sizes) {
size_t total = 0;
for (size_t size : sizes) {
total += size;
}
return total;
}
size_t sum_successful_sizes(const std::vector<int> &results,
const std::vector<size_t> &sizes) {
size_t total = 0;
for (size_t i = 0; i < results.size() && i < sizes.size(); ++i) {
if (results[i] == 0) {
total += sizes[i];
}
}
return total;
}
size_t sum_successful_nested_sizes(
const std::vector<int> &results,
const std::vector<std::vector<size_t>> &nested_sizes) {
size_t total = 0;
for (size_t i = 0; i < results.size() && i < nested_sizes.size(); ++i) {
if (results[i] == 0) {
total += sum_sizes(nested_sizes[i]);
}
}
return total;
}
size_t sum_positive_results(const std::vector<int64_t> &results) {
size_t total = 0;
for (int64_t result : results) {
if (result > 0) {
total += static_cast<size_t>(result);
}
}
return total;
}
size_t sum_positive_results(const std::vector<int> &results) {
size_t total = 0;
for (int result : results) {
if (result > 0) {
total += static_cast<size_t>(result);
}
}
return total;
}
size_t sum_positive_ranges(
const std::vector<std::vector<std::vector<int64_t>>> &results) {
size_t total = 0;
for (const auto &key_rows : results) {
for (const auto &row : key_rows) {
total += sum_positive_results(row);
}
}
return total;
}
size_t sum_buffer_handle_sizes(
const std::vector<std::shared_ptr<BufferHandle>> &buffers) {
size_t total = 0;
for (const auto &buffer : buffers) {
if (buffer != nullptr) {
total += buffer->size();
}
}
return total;
}
PreparedRangedReadRequest prepare_ranged_read_request(
size_t buffer_count, const std::vector<std::vector<std::string>> &all_keys,
const std::vector<std::vector<std::vector<size_t>>> &all_dst_offsets,
@ -535,8 +454,7 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
const std::string &master_server_addr,
const std::shared_ptr<TransferEngine> &transfer_engine,
const std::string &ipc_socket_path, int local_rpc_port,
bool enable_ssd_offload, bool start_offload_rpc_server,
const std::string &ssd_offload_path) {
bool enable_ssd_offload, bool start_offload_rpc_server) {
this->protocol = protocol;
this->ipc_socket_path_ = ipc_socket_path;
const bool should_use_hugepage = use_hugepage_ &&
@ -575,7 +493,7 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
hostname.substr(0, colon_pos + 1) + std::to_string(local_rpc_port);
auto client_opt = mooncake::Client::Create(
this->local_hostname, metadata_server, protocol, device_name,
master_server_addr, transfer_engine, {{"client_mode", "real"}});
master_server_addr, transfer_engine);
if (!client_opt) {
LOG(ERROR) << "Failed to create client";
return tl::unexpected(ErrorCode::INVALID_PARAMS);
@ -604,7 +522,7 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
hostname + ":" + std::to_string(local_rpc_port);
auto client_opt = mooncake::Client::Create(
this->local_hostname, metadata_server, protocol, device_name,
master_server_addr, transfer_engine, {{"client_mode", "real"}});
master_server_addr, transfer_engine);
if (client_opt) {
client_ = *client_opt;
success = true;
@ -681,10 +599,10 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
uint64_t total_glbseg_size = global_segment_size; // For logging
uint64_t current_glbseg_size = 0; // For logging
// For RDMA, auto-discover NUMA nodes with NICs and distribute
// global_segment across them for full NIC utilization.
// In standalone mode with RDMA, auto-discover NUMA nodes with NICs
// and distribute global_segment across them for full NIC utilization.
std::vector<int> seg_numa_nodes;
if (protocol == "rdma") {
if (!ipc_socket_path_.empty() && protocol == "rdma") {
seg_numa_nodes = client_->GetNicNumaNodes();
if (seg_numa_nodes.size() > 1) {
std::string nodes_str;
@ -797,12 +715,8 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
}
if (enable_ssd_offload) {
auto file_storage_config = FileStorageConfig::FromEnvironment();
if (!ssd_offload_path.empty()) {
file_storage_config.storage_filepath = ssd_offload_path;
}
file_storage_ = std::make_shared<FileStorage>(
file_storage_config, client_, this->local_rpc_addr,
client_->GetSsdMetricPtr());
file_storage_config, client_, this->local_rpc_addr);
auto init_result = file_storage_->Init();
if (!init_result) {
LOG(ERROR) << "file storage init failed with error: "
@ -827,12 +741,11 @@ int RealClient::setup_real(
const std::string &protocol, const std::string &rdma_devices,
const std::string &master_server_addr,
const std::shared_ptr<TransferEngine> &transfer_engine,
const std::string &ipc_socket_path, bool enable_ssd_offload,
const std::string &ssd_offload_path) {
const std::string &ipc_socket_path, bool enable_ssd_offload) {
return to_py_ret(setup_internal(
local_hostname, metadata_server, global_segment_size, local_buffer_size,
protocol, rdma_devices, master_server_addr, transfer_engine,
ipc_socket_path, 50052, enable_ssd_offload, true, ssd_offload_path));
ipc_socket_path, 50052, enable_ssd_offload, true));
}
namespace {
@ -923,8 +836,6 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
return tl::unexpected(ErrorCode::INVALID_PARAMS);
}
std::string ssd_offload_path = get_config(config, "ssd_offload_path");
std::string enable_ssd_offload_str =
get_config(config, "enable_ssd_offload", "false");
std::transform(enable_ssd_offload_str.begin(), enable_ssd_offload_str.end(),
@ -936,7 +847,7 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
return setup_internal(local_hostname, metadata_server, global_segment_size,
local_buffer_size, protocol, rdma_devices,
master_server_addr, nullptr, ipc_socket_path, 50052,
enable_ssd_offload, true, ssd_offload_path);
enable_ssd_offload, true);
}
tl::expected<void, ErrorCode> RealClient::initAll_internal(
@ -1165,17 +1076,8 @@ tl::expected<void, ErrorCode> RealClient::put_dummy_helper(
int RealClient::put(const std::string &key, std::span<const char> value,
const ReplicateConfig &config) {
auto result = execute_timed_operation<tl::expected<void, ErrorCode>>(
[&]() {
return put_internal(key, value, config, client_buffer_allocator_);
},
[](const auto &ret) { return ret.has_value(); },
[&](uint64_t latency_us, const auto &) {
client_->ObserveTransferOperation(TransferOperationKind::kWrite,
"put", value.size_bytes(),
latency_us);
});
return to_py_ret(result);
return to_py_ret(
put_internal(key, value, config, client_buffer_allocator_));
}
tl::expected<void, ErrorCode> RealClient::put_batch_internal(
@ -1264,18 +1166,8 @@ tl::expected<void, ErrorCode> RealClient::put_batch_dummy_helper(
int RealClient::put_batch(const std::vector<std::string> &keys,
const std::vector<std::span<const char>> &values,
const ReplicateConfig &config) {
auto result = execute_timed_operation<tl::expected<void, ErrorCode>>(
[&]() {
return put_batch_internal(keys, values, config,
client_buffer_allocator_);
},
[](const auto &ret) { return ret.has_value(); },
[&](uint64_t latency_us, const auto &) {
client_->ObserveTransferOperation(
TransferOperationKind::kWrite, "put_batch",
sum_value_sizes(values), latency_us);
});
return to_py_ret(result);
return to_py_ret(
put_batch_internal(keys, values, config, client_buffer_allocator_));
}
tl::expected<void, ErrorCode> RealClient::put_parts_internal(
@ -1356,18 +1248,8 @@ tl::expected<void, ErrorCode> RealClient::put_parts_dummy_helper(
int RealClient::put_parts(const std::string &key,
std::vector<std::span<const char>> values,
const ReplicateConfig &config) {
auto result = execute_timed_operation<tl::expected<void, ErrorCode>>(
[&]() {
return put_parts_internal(key, values, config,
client_buffer_allocator_);
},
[](const auto &ret) { return ret.has_value(); },
[&](uint64_t latency_us, const auto &) {
client_->ObserveTransferOperation(
TransferOperationKind::kWrite, "put_parts",
sum_value_sizes(values), latency_us);
});
return to_py_ret(result);
return to_py_ret(
put_parts_internal(key, values, config, client_buffer_allocator_));
}
tl::expected<void, ErrorCode> RealClient::remove_internal(
@ -1964,14 +1846,7 @@ std::shared_ptr<BufferHandle> RealClient::get_buffer_internal(
// Implementation of get_buffer method
std::shared_ptr<BufferHandle> RealClient::get_buffer(const std::string &key) {
return execute_timed_operation<std::shared_ptr<BufferHandle>>(
[&]() { return get_buffer_internal(key, client_buffer_allocator_); },
[](const auto &buffer) { return buffer != nullptr; },
[&](uint64_t latency_us, const auto &buffer) {
client_->ObserveTransferOperation(TransferOperationKind::kRead,
"get_buffer", buffer->size(),
latency_us);
});
return get_buffer_internal(key, client_buffer_allocator_);
}
tl::expected<std::tuple<uint64_t, size_t>, ErrorCode>
@ -2261,14 +2136,7 @@ RealClient::batch_get_buffer_internal(
// Implementation of batch_get_buffer method
std::vector<std::shared_ptr<BufferHandle>> RealClient::batch_get_buffer(
const std::vector<std::string> &keys) {
return execute_timed_operation<std::vector<std::shared_ptr<BufferHandle>>>(
[&]() { return batch_get_buffer_internal(keys); },
[](const auto &) { return true; },
[&](uint64_t latency_us, const auto &buffers) {
client_->ObserveTransferOperation(
TransferOperationKind::kRead, "batch_get_buffer",
sum_buffer_handle_sizes(buffers), latency_us);
});
return batch_get_buffer_internal(keys);
}
tl::expected<void, ErrorCode> RealClient::register_buffer_internal(
@ -2423,17 +2291,7 @@ tl::expected<int64_t, ErrorCode> RealClient::get_into_range_internal(
int64_t RealClient::get_into(const std::string &key, void *buffer,
size_t size) {
auto result = execute_timed_operation<tl::expected<int64_t, ErrorCode>>(
[&]() {
return get_into_range_internal(key, buffer, 0, 0, size, true);
},
[](const auto &ret) { return ret.has_value(); },
[&](uint64_t latency_us, const auto &ret) {
client_->ObserveTransferOperation(
TransferOperationKind::kRead, "get_into",
static_cast<uint64_t>(ret.value()), latency_us);
});
return to_py_ret(result);
return to_py_ret(get_into_range_internal(key, buffer, 0, 0, size, true));
}
std::vector<std::vector<std::vector<tl::expected<int64_t, ErrorCode>>>>
@ -2553,20 +2411,8 @@ std::vector<std::vector<std::vector<int64_t>>> RealClient::get_into_ranges(
const std::vector<std::vector<std::vector<size_t>>> &all_dst_offsets,
const std::vector<std::vector<std::vector<size_t>>> &all_src_offsets,
const std::vector<std::vector<std::vector<size_t>>> &all_sizes) {
auto results =
execute_timed_operation<std::vector<std::vector<std::vector<int64_t>>>>(
[&]() {
return convert_ranged_read_results(
get_into_ranges_internal(buffers, all_keys, all_dst_offsets,
all_src_offsets, all_sizes));
},
[](const auto &) { return true; },
[&](uint64_t latency_us, const auto &ret) {
client_->ObserveTransferOperation(
TransferOperationKind::kRead, "get_into_ranges",
sum_positive_ranges(ret), latency_us);
});
return results;
return convert_ranged_read_results(get_into_ranges_internal(
buffers, all_keys, all_dst_offsets, all_src_offsets, all_sizes));
}
std::string RealClient::get_hostname() const { return local_hostname; }
@ -2575,21 +2421,7 @@ std::vector<int> RealClient::batch_put_from(
const std::vector<std::string> &keys, const std::vector<void *> &buffers,
const std::vector<size_t> &sizes, const ReplicateConfig &config) {
auto internal_results =
execute_timed_operation<std::vector<tl::expected<void, ErrorCode>>>(
[&]() {
return batch_put_from_internal(keys, buffers, sizes, config);
},
[](const auto &) { return true; },
[&](uint64_t latency_us, const auto &ret) {
std::vector<int> py_results;
py_results.reserve(ret.size());
for (const auto &item : ret) {
py_results.push_back(to_py_ret(item));
}
client_->ObserveTransferOperation(
TransferOperationKind::kWrite, "batch_put_from",
sum_successful_sizes(py_results, sizes), latency_us);
});
batch_put_from_internal(keys, buffers, sizes, config);
std::vector<int> results;
results.reserve(internal_results.size());
@ -2732,14 +2564,7 @@ tl::expected<void, ErrorCode> RealClient::put_from_internal(
int RealClient::put_from(const std::string &key, void *buffer, size_t size,
const ReplicateConfig &config) {
auto result = execute_timed_operation<tl::expected<void, ErrorCode>>(
[&]() { return put_from_internal(key, buffer, size, config); },
[](const auto &ret) { return ret.has_value(); },
[&](uint64_t latency_us, const auto &) {
client_->ObserveTransferOperation(TransferOperationKind::kWrite,
"put_from", size, latency_us);
});
return to_py_ret(result);
return to_py_ret(put_from_internal(key, buffer, size, config));
}
// --- Upsert implementations ---
@ -2780,18 +2605,8 @@ tl::expected<void, ErrorCode> RealClient::upsert_internal(
int RealClient::upsert(const std::string &key, std::span<const char> value,
const ReplicateConfig &config) {
auto result = execute_timed_operation<tl::expected<void, ErrorCode>>(
[&]() {
return upsert_internal(key, value, config,
client_buffer_allocator_);
},
[](const auto &ret) { return ret.has_value(); },
[&](uint64_t latency_us, const auto &) {
client_->ObserveTransferOperation(TransferOperationKind::kWrite,
"upsert", value.size_bytes(),
latency_us);
});
return to_py_ret(result);
return to_py_ret(
upsert_internal(key, value, config, client_buffer_allocator_));
}
tl::expected<void, ErrorCode> RealClient::upsert_dummy_helper(
@ -2841,14 +2656,7 @@ tl::expected<void, ErrorCode> RealClient::upsert_from_internal(
int RealClient::upsert_from(const std::string &key, void *buffer, size_t size,
const ReplicateConfig &config) {
auto result = execute_timed_operation<tl::expected<void, ErrorCode>>(
[&]() { return upsert_from_internal(key, buffer, size, config); },
[](const auto &ret) { return ret.has_value(); },
[&](uint64_t latency_us, const auto &) {
client_->ObserveTransferOperation(TransferOperationKind::kWrite,
"upsert_from", size, latency_us);
});
return to_py_ret(result);
return to_py_ret(upsert_from_internal(key, buffer, size, config));
}
std::vector<tl::expected<void, ErrorCode>>
@ -2894,21 +2702,7 @@ std::vector<int> RealClient::batch_upsert_from(
const std::vector<std::string> &keys, const std::vector<void *> &buffers,
const std::vector<size_t> &sizes, const ReplicateConfig &config) {
auto internal_results =
execute_timed_operation<std::vector<tl::expected<void, ErrorCode>>>(
[&]() {
return batch_upsert_from_internal(keys, buffers, sizes, config);
},
[](const auto &) { return true; },
[&](uint64_t latency_us, const auto &ret) {
std::vector<int> py_results;
py_results.reserve(ret.size());
for (const auto &item : ret) {
py_results.push_back(to_py_ret(item));
}
client_->ObserveTransferOperation(
TransferOperationKind::kWrite, "batch_upsert_from",
sum_successful_sizes(py_results, sizes), latency_us);
});
batch_upsert_from_internal(keys, buffers, sizes, config);
std::vector<int> results;
results.reserve(internal_results.size());
for (const auto &result : internal_results) {
@ -3032,18 +2826,8 @@ tl::expected<void, ErrorCode> RealClient::upsert_parts_internal(
int RealClient::upsert_parts(const std::string &key,
std::vector<std::span<const char>> values,
const ReplicateConfig &config) {
auto result = execute_timed_operation<tl::expected<void, ErrorCode>>(
[&]() {
return upsert_parts_internal(key, values, config,
client_buffer_allocator_);
},
[](const auto &ret) { return ret.has_value(); },
[&](uint64_t latency_us, const auto &) {
client_->ObserveTransferOperation(
TransferOperationKind::kWrite, "upsert_parts",
sum_value_sizes(values), latency_us);
});
return to_py_ret(result);
return to_py_ret(
upsert_parts_internal(key, values, config, client_buffer_allocator_));
}
tl::expected<void, ErrorCode> RealClient::upsert_parts_dummy_helper(
@ -3146,18 +2930,8 @@ tl::expected<void, ErrorCode> RealClient::upsert_batch_dummy_helper(
int RealClient::upsert_batch(const std::vector<std::string> &keys,
const std::vector<std::span<const char>> &values,
const ReplicateConfig &config) {
auto result = execute_timed_operation<tl::expected<void, ErrorCode>>(
[&]() {
return upsert_batch_internal(keys, values, config,
client_buffer_allocator_);
},
[](const auto &ret) { return ret.has_value(); },
[&](uint64_t latency_us, const auto &) {
client_->ObserveTransferOperation(
TransferOperationKind::kWrite, "upsert_batch",
sum_value_sizes(values), latency_us);
});
return to_py_ret(result);
return to_py_ret(
upsert_batch_internal(keys, values, config, client_buffer_allocator_));
}
// --- End Upsert implementations ---
@ -3165,20 +2939,7 @@ int RealClient::upsert_batch(const std::vector<std::string> &keys,
std::vector<int64_t> RealClient::batch_get_into(
const std::vector<std::string> &keys, const std::vector<void *> &buffers,
const std::vector<size_t> &sizes) {
auto internal_results =
execute_timed_operation<std::vector<tl::expected<int64_t, ErrorCode>>>(
[&]() { return batch_get_into_internal(keys, buffers, sizes); },
[](const auto &) { return true; },
[&](uint64_t latency_us, const auto &ret) {
std::vector<int64_t> py_results;
py_results.reserve(ret.size());
for (const auto &item : ret) {
py_results.push_back(to_py_ret(item));
}
client_->ObserveTransferOperation(
TransferOperationKind::kRead, "batch_get_into",
sum_positive_results(py_results), latency_us);
});
auto internal_results = batch_get_into_internal(keys, buffers, sizes);
std::vector<int64_t> results;
results.reserve(internal_results.size());
@ -3582,7 +3343,6 @@ int RealClient::put_from_with_metadata(const std::string &key, void *buffer,
void *metadata_buffer, size_t size,
size_t metadata_size,
const ReplicateConfig &config) {
const auto start_time = std::chrono::steady_clock::now();
// NOTE: The buffer address must be previously registered with
// register_buffer() for zero-copy RDMA operations to work correctly
if (config.prefer_alloc_in_same_node) {
@ -3625,10 +3385,6 @@ int RealClient::put_from_with_metadata(const std::string &key, void *buffer,
<< toString(put_result.error());
return -toInt(put_result.error());
}
client_->ObserveTransferOperation(
TransferOperationKind::kWrite, "put_from_with_metadata",
size + metadata_size, elapsed_us_since(start_time));
return 0;
}
@ -3637,24 +3393,10 @@ std::vector<int> RealClient::batch_put_from_multi_buffers(
const std::vector<std::vector<void *>> &all_buffers,
const std::vector<std::vector<size_t>> &sizes,
const ReplicateConfig &config) {
auto start = std::chrono::steady_clock::now();
auto internal_results =
execute_timed_operation<std::vector<tl::expected<void, ErrorCode>>>(
[&]() {
return batch_put_from_multi_buffers_internal(keys, all_buffers,
sizes, config);
},
[](const auto &) { return true; },
[&](uint64_t latency_us, const auto &ret) {
std::vector<int> py_results;
py_results.reserve(ret.size());
for (const auto &item : ret) {
py_results.push_back(to_py_ret(item));
}
client_->ObserveTransferOperation(
TransferOperationKind::kWrite,
"batch_put_from_multi_buffers",
sum_successful_nested_sizes(py_results, sizes), latency_us);
});
batch_put_from_multi_buffers_internal(keys, all_buffers, sizes, config);
std::vector<int> results;
results.reserve(internal_results.size());
@ -3662,6 +3404,10 @@ std::vector<int> RealClient::batch_put_from_multi_buffers(
results.push_back(to_py_ret(result));
}
auto duration_call = std::chrono::duration_cast<std::chrono::microseconds>(
std::chrono::steady_clock::now() - start);
VLOG(1) << "batch_put_from_multi_buffers: " << duration_call.count()
<< " us";
return results;
}
@ -3707,30 +3453,19 @@ std::vector<int> RealClient::batch_get_into_multi_buffers(
const std::vector<std::vector<void *>> &all_buffers,
const std::vector<std::vector<size_t>> &all_sizes,
bool prefer_alloc_in_same_node) {
auto internal_results =
execute_timed_operation<std::vector<tl::expected<int64_t, ErrorCode>>>(
[&]() {
return batch_get_into_multi_buffers_internal(
keys, all_buffers, all_sizes, prefer_alloc_in_same_node);
},
[](const auto &) { return true; },
[&](uint64_t latency_us, const auto &ret) {
std::vector<int> py_results;
py_results.reserve(ret.size());
for (const auto &item : ret) {
py_results.push_back(to_py_ret(item));
}
client_->ObserveTransferOperation(
TransferOperationKind::kRead,
"batch_get_into_multi_buffers",
sum_positive_results(py_results), latency_us);
});
auto start = std::chrono::steady_clock::now();
auto internal_results = batch_get_into_multi_buffers_internal(
keys, all_buffers, all_sizes, prefer_alloc_in_same_node);
std::vector<int> results;
results.reserve(internal_results.size());
for (const auto &result : internal_results) {
results.push_back(to_py_ret(result));
}
auto duration_call = std::chrono::duration_cast<std::chrono::microseconds>(
std::chrono::steady_clock::now() - start);
VLOG(1) << "batch_get_into_multi_buffers: " << duration_call.count()
<< " us";
return results;
}
@ -4156,21 +3891,6 @@ RealClient::batch_get_replica_desc(const std::vector<std::string> &keys) {
return replica_map;
}
std::vector<std::string> RealClient::batch_replica_clear(
const std::vector<std::string> &keys, const std::string &segment_name) {
if (!client_) {
LOG(ERROR) << "batch_replica_clear: client not initialized";
return {};
}
auto result =
client_->BatchReplicaClear(keys, client_->getClientId(), segment_name);
if (result) {
return result.value();
}
LOG(ERROR) << "batch_replica_clear failed: " << toString(result.error());
return {};
}
tl::expected<UUID, ErrorCode> RealClient::create_copy_task(
const std::string &key, const std::vector<std::string> &targets) {
return client_->CreateCopyTask(key, targets);
@ -4266,18 +3986,6 @@ ClientRequester::ClientRequester() {
pool_conf.client_config.socket_config =
coro_io::ib_socket_t::config_t{};
}
// Configure reasonable retry limits for SSD offload RPC connections.
// - connect_retry_count: Maximum connection retry attempts (default: 3)
// - reconnect_wait_time: Wait time between retries (default: 1000ms)
// - host_alive_detect_duration: Duration for background alive detection.
// Set to 0 to disable infinite background reconnection attempts when
// a Store node goes down. This prevents continuous "Connection refused"
// logs. When Master cleans up stale local_disk replicas (via
// CleanupStaleHandles), new requests won't route to dead nodes anyway.
pool_conf.connect_retry_count = 3;
pool_conf.reconnect_wait_time = std::chrono::milliseconds{1000};
pool_conf.host_alive_detect_duration = std::chrono::milliseconds{0};
client_pools_ =
std::make_shared<coro_io::client_pools<coro_rpc::coro_rpc_client>>(
pool_conf);

View File

@ -91,10 +91,6 @@ int main(int argc, char *argv[]) {
mooncake::ResourceTracker::getInstance();
gflags::ParseCommandLineFlags(&argc, &argv, true);
if (!FLAGS_log_dir.empty()) {
google::InitGoogleLogging(argv[0]);
}
size_t global_segment_size = string_to_byte_size(FLAGS_global_segment_size);
#ifdef USE_ASCEND_DIRECT
// just set to true, does not affect GPU process.
@ -117,11 +113,10 @@ int main(int argc, char *argv[]) {
return -1;
}
coro_rpc::coro_rpc_server server(FLAGS_threads, FLAGS_port, FLAGS_host);
coro_rpc::coro_rpc_server server(FLAGS_threads, FLAGS_port, "127.0.0.1");
RegisterClientRpcService(server, *client_inst);
LOG(INFO) << "Starting real client service on " << FLAGS_host << ":"
<< FLAGS_port;
LOG(INFO) << "Starting real client service on 127.0.0.1:" << FLAGS_port;
return server.start();
}

View File

@ -273,15 +273,6 @@ ErrorCode ScopedSegmentAccess::GetClientSegments(
return ErrorCode::OK;
}
void ScopedSegmentAccess::UnmountLocalDiskSegment(const UUID& client_id) {
auto it = segment_manager_->client_local_disk_segment_.find(client_id);
if (it != segment_manager_->client_local_disk_segment_.end()) {
segment_manager_->client_local_disk_segment_.erase(it);
LOG(INFO) << "client_id=" << client_id
<< ", action=unmount_local_disk_segment";
}
}
ErrorCode ScopedSegmentAccess::GetAllSegments(
std::vector<std::string>& all_segments) {
all_segments.clear();

View File

@ -49,9 +49,8 @@ FilePerKeyConfig FilePerKeyConfig::FromEnvironment() {
config.fsdir = GetEnvStringOr("MOONCAKE_OFFLOAD_FSDIR", config.fsdir);
config.enable_eviction = GetEnvOr<bool>(
"MOONCAKE_OFFLOAD_ENABLE_EVICTION",
GetEnvOr<bool>("ENABLE_EVICTION", config.enable_eviction));
config.enable_eviction =
GetEnvOr<bool>("ENABLE_EVICTION", config.enable_eviction);
return config;
}
@ -65,14 +64,11 @@ BucketBackendConfig BucketBackendConfig::FromEnvironment() {
config.bucket_size_limit = GetEnvOr<int64_t>(
"MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES", config.bucket_size_limit);
config.max_total_size =
GetEnvOr<int64_t>("MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE",
GetEnvOr<int64_t>("MOONCAKE_BUCKET_MAX_TOTAL_SIZE",
config.max_total_size));
config.max_total_size = GetEnvOr<int64_t>("MOONCAKE_BUCKET_MAX_TOTAL_SIZE",
config.max_total_size);
const auto policy_str = GetEnvStringOr(
"MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY",
GetEnvStringOr("MOONCAKE_BUCKET_EVICTION_POLICY", "none"));
const auto policy_str =
GetEnvStringOr("MOONCAKE_BUCKET_EVICTION_POLICY", "none");
if (policy_str == "fifo") {
config.eviction_policy = BucketEvictionPolicy::FIFO;
} else if (policy_str == "lru") {
@ -1722,18 +1718,6 @@ tl::expected<void, ErrorCode> BucketStorageBackend::Init() {
<< orphaned_space_freed << " bytes";
}
// When max_total_size is not explicitly set (<= 0), default to 90% of
// the physical disk capacity to match FilePerKey backend behavior.
if (bucket_backend_config_.max_total_size <= 0) {
constexpr double kDefaultQuotaPercentage = 0.9;
const auto space_info = fs::space(storage_path_);
bucket_backend_config_.max_total_size = static_cast<int64_t>(
space_info.capacity * kDefaultQuotaPercentage);
LOG(INFO) << "Bucket backend max_total_size not set; using "
<< kDefaultQuotaPercentage * 100 << "% of disk capacity: "
<< bucket_backend_config_.max_total_size << " bytes";
}
bucket_id_generator_.emplace(max_bucket_id);
if (max_bucket_id == BucketIdGenerator::INIT_NEW_START_ID) {
LOG(INFO) << "Initialized BucketIdGenerator with fresh start. "
@ -2203,7 +2187,8 @@ BucketStorageBackend::PendingEviction BucketStorageBackend::PrepareEviction(
int64_t required_size) {
PendingEviction result;
if (bucket_backend_config_.eviction_policy == BucketEvictionPolicy::NONE) {
if (bucket_backend_config_.eviction_policy == BucketEvictionPolicy::NONE ||
bucket_backend_config_.max_total_size <= 0) {
return result;
}

View File

@ -4,8 +4,6 @@
#include <algorithm>
#include <cstdlib>
#include <sstream>
#include <vector>
#include "transfer_engine.h"
#include "transport/transport.h"
@ -225,13 +223,9 @@ bool TransferEngineOperationState::is_completed() {
}
void TransferEngineOperationState::check_task_status() {
// Check all transfers in the batch.
// Wait for ALL tasks to reach a terminal state before setting the result,
// even if some have already failed. This prevents the caller from seeing
// "completed" while background transfers are still in progress, which
// could cause issues when freeBatchID is called in the destructor.
bool all_terminated = true;
std::vector<size_t> failed_task_ids;
// Check all transfers in the batch
bool all_completed = true;
bool has_failure = false;
for (size_t i = 0; i < batch_size_; ++i) {
TransferStatus status;
@ -246,45 +240,38 @@ void TransferEngineOperationState::check_task_status() {
switch (status.s) {
case TransferStatusEnum::COMPLETED:
// This transfer is done successfully
// This transfer is done, continue checking others
break;
case TransferStatusEnum::FAILED:
case TransferStatusEnum::CANCELED:
case TransferStatusEnum::INVALID:
#ifndef USE_ASCEND_DIRECT
VLOG(1) << "Transfer failed for batch " << batch_id_ << " task "
<< i << " with status " << static_cast<int>(status.s);
LOG(ERROR) << "Transfer failed for batch " << batch_id_
<< " task " << i << " with status "
<< static_cast<int>(status.s);
#endif
failed_task_ids.push_back(i);
has_failure = true;
break;
default:
// Transfer is still in progress (WAITING, PENDING, etc.)
all_terminated = false;
// Transfer is still pending (PENDING, RUNNING, etc.)
all_completed = false;
break;
}
}
if (!all_terminated) {
// Some tasks are still in progress; wait for next poll iteration.
// Do NOT set result yet, even if some tasks have already failed.
if (has_failure) {
VLOG(1) << "Setting batch " << batch_id_
<< " result to TRANSFER_FAIL due to task failures";
set_result_internal(ErrorCode::TRANSFER_FAIL);
return;
}
// All tasks have reached a terminal state.
ErrorCode ec = ErrorCode::OK;
if (!failed_task_ids.empty()) {
std::ostringstream oss;
for (size_t j = 0; j < failed_task_ids.size(); ++j) {
if (j > 0) oss << ", ";
oss << failed_task_ids[j];
}
LOG(ERROR) << "Batch " << batch_id_
<< " completed with task failures: task_ids=[" << oss.str()
<< "]";
ec = ErrorCode::TRANSFER_FAIL;
if (all_completed) {
set_result_internal(ErrorCode::OK);
return;
}
set_result_internal(ec);
return;
}
void TransferEngineOperationState::set_result_internal(ErrorCode error_code) {
@ -307,8 +294,7 @@ void TransferEngineOperationState::wait_for_completion() {
return;
}
// 60 seconds
constexpr int64_t timeout_milliseconds = 60 * 1000;
constexpr int64_t timeout_seconds = 60;
#ifdef USE_EVENT_DRIVEN_COMPLETION
VLOG(1) << "Waiting for transfer engine completion for batch " << batch_id_;
@ -328,18 +314,10 @@ void TransferEngineOperationState::wait_for_completion() {
// lock. Under the mutex, relaxed is sufficient; the mutex acquire
// orders prior writes.
std::unique_lock<std::mutex> lock(batch_desc.completion_mutex);
const int64_t elapsed_milliseconds =
getCurrentTimeInMilli() - start_ts_;
if (elapsed_milliseconds < timeout_milliseconds) {
completed = batch_desc.completion_cv.wait_for(
lock,
std::chrono::milliseconds(timeout_milliseconds -
elapsed_milliseconds),
[&batch_desc] {
return batch_desc.is_finished.load(
std::memory_order_relaxed);
});
}
completed = batch_desc.completion_cv.wait_for(
lock, std::chrono::seconds(timeout_seconds), [&batch_desc] {
return batch_desc.is_finished.load(std::memory_order_relaxed);
});
} // Explicitly release completion_mutex before acquiring mutex_
// Once completion is observed, read failure flag.
@ -360,18 +338,20 @@ void TransferEngineOperationState::wait_for_completion() {
VLOG(1) << "Transfer engine operation completed for batch " << batch_id_
<< " with result: " << static_cast<int>(error_code);
} else {
LOG(ERROR) << "Failed to complete transfers after "
<< timeout_milliseconds << " milliseconds for batch "
<< batch_id_;
LOG(ERROR) << "Failed to complete transfers after " << timeout_seconds
<< " seconds for batch " << batch_id_;
}
#else
VLOG(1) << "Starting transfer engine polling for batch " << batch_id_;
constexpr int64_t kOneSecondInNano = 1000 * 1000 * 1000;
const int64_t start_ts = getCurrentTimeInNano();
while (true) {
if (getCurrentTimeInMilli() - start_ts_ > timeout_milliseconds) {
if (getCurrentTimeInNano() - start_ts >
timeout_seconds * kOneSecondInNano) {
LOG(ERROR) << "Failed to complete transfers after "
<< timeout_milliseconds << " milliseconds for batch "
<< batch_id_;
<< timeout_seconds << " seconds for batch " << batch_id_;
set_result_internal(ErrorCode::TRANSFER_FAIL);
return;
}
@ -429,17 +409,10 @@ TransferSubmitter::TransferSubmitter(TransferEngine& engine,
memcpy_pool_(std::make_unique<MemcpyWorkerPool>()),
fileread_pool_(std::make_unique<FilereadWorkerPool>(backend)),
transfer_metric_(transfer_metric) {
// Read MC_STORE_MEMCPY environment variable.
// When not set, auto-detect based on transport type:
// - TCP-only environment: enable memcpy (avoids TCP loopback overhead)
// - RDMA/other transports: disable memcpy (RDMA is more efficient)
// Read MC_STORE_MEMCPY environment variable, default to false (disabled)
const char* env_value = std::getenv("MC_STORE_MEMCPY");
if (env_value == nullptr) {
memcpy_enabled_ = engine_.isTcpOnly();
LOG(INFO) << "MC_STORE_MEMCPY not set, auto-detected: "
<< (memcpy_enabled_ ? "TCP-only environment, memcpy enabled"
: "non-TCP transport available, memcpy "
"disabled");
memcpy_enabled_ = false; // Default: disabled
} else {
std::string env_str(env_value);
// Convert to lowercase for case-insensitive comparison

View File

@ -39,7 +39,6 @@ add_store_test(eviction_strategy_test eviction_strategy_test.cpp)
add_store_test(master_service_test master_service_test.cpp)
add_store_test(batch_remove_test batch_remove_test.cpp)
add_store_test(master_service_ssd_test master_service_ssd_test.cpp)
add_store_test(offload_on_evict_test offload_on_evict_test.cpp)
add_store_test(master_service_ssd_test_for_snapshot
ha/snapshot/master_service_ssd_test_for_snapshot.cpp)
add_store_test(client_integration_test client_integration_test.cpp)
@ -58,7 +57,6 @@ add_store_test(client_local_hot_cache_test client_local_hot_cache_test.cpp)
add_store_test(pybind_client_test pybind_client_test.cpp)
add_store_test(ipv6_client_test ipv6_client_test.cpp)
add_store_test(client_metrics_test client_metrics_test.cpp)
add_store_test(ssd_metrics_test ssd_metrics_test.cpp)
add_store_test(serializer_test serializer_test.cpp)
add_store_test(embedded_snapshot_catalog_store_test
ha/snapshot/catalog/backends/embedded/embedded_snapshot_catalog_store_test.cpp)

View File

@ -45,8 +45,6 @@ TEST_F(ClientMetricsTest, TransferMetricsSummaryTest) {
// Check byte formatting
EXPECT_TRUE(summary.find("Total Read: 1.00 KB") != std::string::npos);
EXPECT_TRUE(summary.find("Total Write: 2.00 MB") != std::string::npos);
EXPECT_TRUE(summary.find("Average Read Throughput:") != std::string::npos);
EXPECT_TRUE(summary.find("Average Write Throughput:") != std::string::npos);
// Check latency summaries
EXPECT_TRUE(summary.find("Get: count=3") != std::string::npos);
@ -111,23 +109,15 @@ TEST_F(ClientMetricsTest, ClientMetricsSummaryTest) {
std::array<std::string, 1> exist_key_label = {"ExistKey"};
metrics.master_client_metric.rpc_count.inc(exist_key_label);
metrics.master_client_metric.rpc_latency.observe(exist_key_label, 180);
metrics.ObserveTransferOperation(TransferOperationKind::kRead, "get_buffer",
2 * 1024, 220);
metrics.ObserveTransferOperation(TransferOperationKind::kWrite, "put_batch",
4 * 1024, 420);
std::string summary = metrics.summary_metrics();
// Should contain transfer, RPC, and interface metrics
// Should contain both transfer and RPC metrics
EXPECT_TRUE(summary.find("Transfer Metrics Summary") != std::string::npos);
EXPECT_TRUE(summary.find("RPC Metrics Summary") != std::string::npos);
EXPECT_TRUE(summary.find("Interface Operation Metrics Summary") !=
std::string::npos);
EXPECT_TRUE(summary.find("Total Read: 5.00 MB") != std::string::npos);
EXPECT_TRUE(summary.find("Total Write: 10.00 MB") != std::string::npos);
EXPECT_TRUE(summary.find("ExistKey: count=1") != std::string::npos);
EXPECT_TRUE(summary.find("get_buffer: count=1") != std::string::npos);
EXPECT_TRUE(summary.find("put_batch: count=1") != std::string::npos);
std::cout << "Full Client Metrics Summary:\n" << summary << std::endl;
}
@ -185,33 +175,6 @@ TEST_F(ClientMetricsTest, CompareWithSerializedMetrics) {
summary.find("No data") != std::string::npos);
}
TEST_F(ClientMetricsTest, BandwidthSummaryRespectsEnvFlag) {
setenv("MC_STORE_CLIENT_METRIC_BANDWIDTH", "0", 1);
auto metrics = ClientMetric::Create();
ASSERT_NE(metrics, nullptr);
metrics->transfer_metric.total_read_bytes.inc(1024);
std::string summary = metrics->summary_metrics();
EXPECT_TRUE(summary.find("Average Read Throughput:") == std::string::npos);
unsetenv("MC_STORE_CLIENT_METRIC_BANDWIDTH");
}
TEST_F(ClientMetricsTest, SummaryCanOmitMasterRpcMetrics) {
auto metrics = ClientMetric::Create({}, false);
ASSERT_NE(metrics, nullptr);
metrics->ObserveTransferOperation(TransferOperationKind::kRead,
"get_buffer", 1024, 200);
std::string summary = metrics->summary_metrics();
std::string serialized;
metrics->serialize(serialized);
EXPECT_TRUE(summary.find("RPC Metrics Summary") == std::string::npos);
EXPECT_TRUE(serialized.find("mooncake_client_rpc_count") ==
std::string::npos);
}
TEST_F(ClientMetricsTest, SerializeWithDynamicLabels) {
auto verify = [](const std::string& str) {
EXPECT_TRUE(str.find("instance_id=\"12345\"") != std::string::npos);

View File

@ -8,7 +8,6 @@
#include "storage_backend.h"
#include "file_storage.h"
#include "utils/common.h"
#include "client_metric.h"
namespace mooncake {
@ -26,8 +25,7 @@ class FileStorageTest : public ::testing::Test {
FLAGS_logtostderr = true;
UnsetEnv("MOONCAKE_OFFLOAD_FILE_STORAGE_PATH");
UnsetEnv("MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES");
UnsetEnv("MOONCAKE_OFFLOAD_SCANMETA_ITERATOR_KEYS_LIMIT");
UnsetEnv("MOONCAKE_SCANMETA_ITERATOR_KEYS_LIMIT");
UnsetEnv("MOONCAKE_OFFLOAD_BUCKET_ITERATOR_KEYS_LIMIT");
UnsetEnv("MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT");
UnsetEnv("MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES");
UnsetEnv("MOONCAKE_OFFLOAD_TOTAL_KEYS_LIMIT");
@ -428,118 +426,4 @@ TEST_F(FileStorageTest, BatchLoad_WithStorageBackendAdaptor) {
}
}
TEST_F(FileStorageTest, BatchLoadRecordsSsdMetrics) {
// Setup: write data to storage backend via BatchOffloadUtil
std::vector<std::string> keys;
std::vector<int64_t> sizes;
std::unordered_map<std::string, std::string> batch_data;
auto file_storage_config = FileStorageConfig::FromEnvironment();
file_storage_config.storage_filepath = data_path;
// Create FileStorage WITH SsdMetric
SsdMetric ssd_metric;
FileStorage fileStorage(file_storage_config, nullptr, "localhost:9003",
&ssd_metric);
// Write test data to disk
ASSERT_TRUE(FileStorageBatchOffload(fileStorage, keys, sizes, batch_data));
ASSERT_FALSE(keys.empty());
// Allocate buffers and call BatchLoad (read path)
auto allocate_res = FileStorageAllocateBatch(fileStorage, keys, sizes);
ASSERT_TRUE(allocate_res);
auto load_result =
FileStorageBatchLoad(fileStorage, allocate_res.value()->slices);
ASSERT_TRUE(load_result);
// Verify SSD read metrics were recorded
EXPECT_EQ(ssd_metric.ssd_read_ops.value(),
static_cast<int64_t>(allocate_res.value()->slices.size()));
// Verify bytes: sum of all slice sizes
int64_t expected_bytes = 0;
for (const auto& [key, slice] : allocate_res.value()->slices) {
expected_bytes += slice.size;
}
EXPECT_EQ(ssd_metric.ssd_read_bytes.value(), expected_bytes);
EXPECT_GT(expected_bytes, 0);
// Verify latency histogram has exactly 1 observation (one BatchLoad call)
auto buckets = ssd_metric.ssd_read_latency_us.get_bucket_counts();
int64_t total_observations = 0;
for (auto& b : buckets) {
total_observations += b->value();
}
EXPECT_EQ(total_observations, 1);
// Write metrics should remain 0 (BatchOffloadUtil bypasses FileStorage)
EXPECT_EQ(ssd_metric.ssd_write_ops.value(), 0);
EXPECT_EQ(ssd_metric.ssd_write_bytes.value(), 0);
LOG(INFO) << "SSD read metrics after BatchLoad: ops="
<< ssd_metric.ssd_read_ops.value()
<< ", bytes=" << ssd_metric.ssd_read_bytes.value();
}
TEST_F(FileStorageTest, BatchLoadFailureDoesNotRecordSsdMetrics) {
auto file_storage_config = FileStorageConfig::FromEnvironment();
file_storage_config.storage_filepath = data_path;
SsdMetric ssd_metric;
FileStorage fileStorage(file_storage_config, nullptr, "localhost:9003",
&ssd_metric);
// Init storage backend so we can call BatchLoad
// But load with keys that don't exist on disk -> should fail
std::unordered_map<std::string, Slice> batch_object;
char dummy_buf[4096] = {};
batch_object["nonexistent_key_1"] = Slice{dummy_buf, 4096};
batch_object["nonexistent_key_2"] = Slice{dummy_buf, 8192};
auto result = FileStorageBatchLoad(fileStorage, batch_object);
// Expect failure (keys don't exist on disk)
EXPECT_FALSE(result);
// Metrics should remain 0 - failed operations are not counted
EXPECT_EQ(ssd_metric.ssd_read_ops.value(), 0);
EXPECT_EQ(ssd_metric.ssd_read_bytes.value(), 0);
auto buckets = ssd_metric.ssd_read_latency_us.get_bucket_counts();
int64_t total_observations = 0;
for (auto& b : buckets) {
total_observations += b->value();
}
EXPECT_EQ(total_observations, 0);
LOG(INFO) << "SSD metrics after failed BatchLoad: ops="
<< ssd_metric.ssd_read_ops.value()
<< ", bytes=" << ssd_metric.ssd_read_bytes.value();
}
TEST_F(FileStorageTest, NullSsdMetricDoesNotCrash) {
// FileStorage with nullptr SsdMetric should work fine (no metrics recorded)
std::vector<std::string> keys;
std::vector<int64_t> sizes;
std::unordered_map<std::string, std::string> batch_data;
auto file_storage_config = FileStorageConfig::FromEnvironment();
file_storage_config.storage_filepath = data_path;
// nullptr SsdMetric (default)
FileStorage fileStorage(file_storage_config, nullptr, "localhost:9003");
ASSERT_TRUE(FileStorageBatchOffload(fileStorage, keys, sizes, batch_data));
ASSERT_FALSE(keys.empty());
auto allocate_res = FileStorageAllocateBatch(fileStorage, keys, sizes);
ASSERT_TRUE(allocate_res);
auto load_result =
FileStorageBatchLoad(fileStorage, allocate_res.value()->slices);
ASSERT_TRUE(load_result);
// No crash = success. No metrics pointer, so nothing to verify.
}
} // namespace mooncake

View File

@ -80,20 +80,6 @@ class HealthCheckTest : public ::testing::Test {
16 * 1024 * 1024, FLAGS_protocol,
rdma_devices, master_address_);
}
bool WaitForHealthCode(
int expected_code,
std::chrono::milliseconds timeout = std::chrono::seconds(10),
std::chrono::milliseconds interval = std::chrono::milliseconds(100)) {
const auto deadline = std::chrono::steady_clock::now() + timeout;
while (std::chrono::steady_clock::now() < deadline) {
if (py_client_->health_check() == expected_code) {
return true;
}
std::this_thread::sleep_for(interval);
}
return py_client_->health_check() == expected_code;
}
};
// Test 1: health_check returns HC_NOT_INITIALIZED before setup
@ -134,8 +120,9 @@ TEST_F(HealthCheckTest, ReturnsTwoWhenMasterDown) {
// Stop master, wait for ping to fail
master_.Stop();
EXPECT_TRUE(WaitForHealthCode(HC_MASTER_UNREACHABLE))
<< "Timed out waiting for HC_MASTER_UNREACHABLE";
std::this_thread::sleep_for(std::chrono::seconds(3));
EXPECT_EQ(py_client_->health_check(), HC_MASTER_UNREACHABLE);
py_client_->tearDownAll();
}
@ -167,8 +154,7 @@ TEST_F(HealthCheckTest, HttpReturns503WhenMasterDown) {
EXPECT_EQ(py_client_->health_check(), HC_HEALTHY);
master_.Stop();
ASSERT_TRUE(WaitForHealthCode(HC_MASTER_UNREACHABLE))
<< "Timed out waiting for HC_MASTER_UNREACHABLE";
std::this_thread::sleep_for(std::chrono::seconds(3));
auto resp = fetch_health(http_port);
EXPECT_EQ(resp.http_status, 503);

View File

@ -1,388 +0,0 @@
#include "master_service.h"
#include <glog/logging.h>
#include <gtest/gtest.h>
#include <atomic>
#include <chrono>
#include <memory>
#include <string>
#include <thread>
#include <vector>
#include "types.h"
namespace mooncake::test {
class OffloadOnEvictTest : public ::testing::Test {
protected:
void SetUp() override {
google::InitGoogleLogging("OffloadOnEvictTest");
FLAGS_logtostderr = true;
}
void TearDown() override { google::ShutdownGoogleLogging(); }
static constexpr size_t kDefaultSegmentBase = 0x300000000;
Segment MakeSegment(std::string name, size_t base, size_t size) const {
Segment segment;
segment.id = generate_uuid();
segment.name = std::move(name);
segment.base = base;
segment.size = size;
segment.te_endpoint = segment.name;
return segment;
}
struct MountedSegmentContext {
UUID segment_id;
UUID client_id;
};
MountedSegmentContext PrepareSegment(MasterService& service,
std::string name, size_t base,
size_t size) const {
Segment segment = MakeSegment(std::move(name), base, size);
UUID client_id = generate_uuid();
auto mount_result = service.MountSegment(segment, client_id);
EXPECT_TRUE(mount_result.has_value());
return {.segment_id = segment.id, .client_id = client_id};
}
// Put an object and complete it.
void PutObject(MasterService& service, const UUID& client_id,
const std::string& key, size_t size = 1024) {
ReplicateConfig config;
config.replica_num = 1;
auto put_start = service.PutStart(client_id, key, size, config);
ASSERT_TRUE(put_start.has_value()) << "PutStart failed for key=" << key;
auto put_end = service.PutEnd(client_id, key, ReplicaType::MEMORY);
ASSERT_TRUE(put_end.has_value()) << "PutEnd failed for key=" << key;
}
// Drain the offload queue via OffloadObjectHeartbeat.
std::unordered_map<std::string, int64_t> DrainOffloadQueue(
MasterService& service, const UUID& client_id) {
auto res = service.OffloadObjectHeartbeat(client_id, true);
if (!res) {
return {};
}
return std::move(res.value());
}
template <typename Predicate>
void WaitUntil(
Predicate&& predicate,
std::chrono::milliseconds timeout = std::chrono::milliseconds(4000),
std::chrono::milliseconds interval =
std::chrono::milliseconds(50)) const {
const auto deadline = std::chrono::steady_clock::now() + timeout;
while (std::chrono::steady_clock::now() < deadline) {
if (predicate()) {
return;
}
std::this_thread::sleep_for(interval);
}
EXPECT_TRUE(predicate());
}
// Fill a segment until PutStart fails, triggering eviction.
// Returns the number of successful puts.
int FillSegmentUntilEviction(MasterService& service, const UUID& client_id,
const std::string& key_prefix,
size_t object_size, int max_puts) {
int success_puts = 0;
for (int i = 0; i < max_puts; ++i) {
std::string key = key_prefix + std::to_string(i);
ReplicateConfig config;
config.replica_num = 1;
auto result = service.PutStart(client_id, key, object_size, config);
if (result.has_value()) {
auto end = service.PutEnd(client_id, key, ReplicaType::MEMORY);
EXPECT_TRUE(end.has_value());
success_puts++;
} else {
// Wait for eviction to process
std::this_thread::sleep_for(std::chrono::milliseconds(50));
}
}
return success_puts;
}
};
// =============================================================================
// Combo A: Default config (offload at PutEnd)
// =============================================================================
TEST_F(OffloadOnEvictTest, ComboA_OffloadAtPutEnd) {
MasterServiceConfig config;
config.enable_offload = true;
config.default_kv_lease_ttl = 2000;
auto service = std::make_unique<MasterService>(config);
constexpr size_t seg_size = 1024 * 1024 * 16;
auto ctx =
PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size);
// Mount local disk segment with offloading ENABLED
auto mount_ld = service->MountLocalDiskSegment(ctx.client_id, true);
ASSERT_TRUE(mount_ld.has_value());
// Put objects
PutObject(*service, ctx.client_id, "key_a1");
PutObject(*service, ctx.client_id, "key_a2");
PutObject(*service, ctx.client_id, "key_a3");
// Default mode: PutEnd pushes to offload queue immediately
auto queued = DrainOffloadQueue(*service, ctx.client_id);
EXPECT_EQ(queued.size(), 3u)
<< "Default: all 3 objects should be in offload queue after PutEnd";
EXPECT_TRUE(queued.count("key_a1"));
EXPECT_TRUE(queued.count("key_a2"));
EXPECT_TRUE(queued.count("key_a3"));
service->RemoveAll();
}
TEST_F(OffloadOnEvictTest, ComboA_EvictionWorks) {
// Regression: eviction still works in default mode
const uint64_t kv_lease_ttl = 2000;
MasterServiceConfig config;
config.enable_offload = true;
config.default_kv_lease_ttl = kv_lease_ttl;
auto service = std::make_unique<MasterService>(config);
// Large segment: can hold ~16K objects of 15KB
constexpr size_t seg_size = 1024 * 1024 * 16 * 15;
constexpr size_t object_size = 1024 * 15;
auto ctx =
PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size);
// Put more objects than the segment can hold
int success_puts = FillSegmentUntilEviction(
*service, ctx.client_id, "evict_a_", object_size, 1024 * 16 + 50);
EXPECT_GT(success_puts, 1024 * 16)
<< "Default: eviction should allow more puts than capacity";
std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
service->RemoveAll();
}
// =============================================================================
// Combo B: offload_on_evict=true (offload on evict, no force-evict)
// =============================================================================
TEST_F(OffloadOnEvictTest, ComboB_PutEndSkipsOffloadQueue) {
MasterServiceConfig config;
config.enable_offload = true;
config.offload_on_evict = true;
config.default_kv_lease_ttl = 2000;
auto service = std::make_unique<MasterService>(config);
constexpr size_t seg_size = 1024 * 1024 * 16;
auto ctx =
PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size);
auto mount_ld = service->MountLocalDiskSegment(ctx.client_id, true);
ASSERT_TRUE(mount_ld.has_value());
PutObject(*service, ctx.client_id, "key_b1");
PutObject(*service, ctx.client_id, "key_b2");
PutObject(*service, ctx.client_id, "key_b3");
// Offload-on-evict: PutEnd should NOT push to offload queue
auto queued = DrainOffloadQueue(*service, ctx.client_id);
EXPECT_EQ(queued.size(), 0u)
<< "Offload-on-evict: queue should be empty after PutEnd";
service->RemoveAll();
}
TEST_F(OffloadOnEvictTest, ComboB_EvictionTriggersOffload) {
const uint64_t kv_lease_ttl = 2000;
MasterServiceConfig config;
config.enable_offload = true;
config.offload_on_evict = true;
config.default_kv_lease_ttl = kv_lease_ttl;
auto service = std::make_unique<MasterService>(config);
constexpr size_t seg_size = 1024 * 1024 * 16 * 15;
constexpr size_t object_size = 1024 * 15;
auto ctx =
PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size);
auto mount_ld = service->MountLocalDiskSegment(ctx.client_id, true);
ASSERT_TRUE(mount_ld.has_value());
// Fill segment to trigger eviction
bool eviction_triggered = false;
int success_puts = 0;
for (int i = 0; i < 1024 * 16 + 50; ++i) {
std::string key = "evict_b_" + std::to_string(i);
ReplicateConfig config;
config.replica_num = 1;
auto result =
service->PutStart(ctx.client_id, key, object_size, config);
if (result.has_value()) {
auto end = service->PutEnd(ctx.client_id, key, ReplicaType::MEMORY);
ASSERT_TRUE(end.has_value());
success_puts++;
} else {
eviction_triggered = true;
std::this_thread::sleep_for(std::chrono::milliseconds(50));
}
}
EXPECT_TRUE(eviction_triggered)
<< "Eviction should trigger when segment fills up";
// Offload-on-evict: eviction should push objects to offload queue
auto queued = DrainOffloadQueue(*service, ctx.client_id);
EXPECT_GT(queued.size(), 0u)
<< "Offload-on-evict: eviction should push to offload queue";
std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
service->RemoveAll();
}
TEST_F(OffloadOnEvictTest, ComboB_NoFallbackWithoutForceEvict) {
// Without force_evict AND without a LocalDiskSegment, offload queue push
// fails and eviction does NOT force-delete MEMORY (data-preserving).
// The segment fills and subsequent puts fail — this is the safe default.
const uint64_t kv_lease_ttl = 2000;
MasterServiceConfig config;
config.enable_offload = true;
config.offload_on_evict = true;
config.default_kv_lease_ttl = kv_lease_ttl;
auto service = std::make_unique<MasterService>(config);
// NO local disk segment mounted — PushOffloadingQueue will fail
constexpr size_t seg_size = 1024 * 1024 * 16 * 15;
constexpr size_t object_size = 1024 * 15;
auto ctx =
PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size);
// Without force_evict, push failures mean DRAM cannot be freed,
// so we can only put up to segment capacity (no overflow).
int success_puts = FillSegmentUntilEviction(
*service, ctx.client_id, "evict_b2_", object_size, 1024 * 16 + 50);
EXPECT_LE(success_puts, 1024 * 16)
<< "Without force_evict, segment should fill and stay full";
std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
service->RemoveAll();
}
// =============================================================================
// Combo C: offload_on_evict=true + offload_force_evict=true
// =============================================================================
TEST_F(OffloadOnEvictTest, ComboC_PutEndSkipsOffloadQueue) {
MasterServiceConfig config;
config.enable_offload = true;
config.offload_on_evict = true;
config.offload_force_evict = true;
config.default_kv_lease_ttl = 2000;
auto service = std::make_unique<MasterService>(config);
constexpr size_t seg_size = 1024 * 1024 * 16;
auto ctx =
PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size);
auto mount_ld = service->MountLocalDiskSegment(ctx.client_id, true);
ASSERT_TRUE(mount_ld.has_value());
PutObject(*service, ctx.client_id, "key_c1");
PutObject(*service, ctx.client_id, "key_c2");
// Same as Combo B: PutEnd should skip offload queue
auto queued = DrainOffloadQueue(*service, ctx.client_id);
EXPECT_EQ(queued.size(), 0u)
<< "Combo C: offload queue should be empty after PutEnd";
service->RemoveAll();
}
TEST_F(OffloadOnEvictTest, ComboC_EvictionWithForceEvict) {
const uint64_t kv_lease_ttl = 2000;
MasterServiceConfig config;
config.enable_offload = true;
config.offload_on_evict = true;
config.offload_force_evict = true;
config.default_kv_lease_ttl = kv_lease_ttl;
auto service = std::make_unique<MasterService>(config);
constexpr size_t seg_size = 1024 * 1024 * 16 * 15;
constexpr size_t object_size = 1024 * 15;
auto ctx =
PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size);
auto mount_ld = service->MountLocalDiskSegment(ctx.client_id, true);
ASSERT_TRUE(mount_ld.has_value());
// With force-evict, eviction should work effectively.
// Note: without a real FileStorage heartbeat, offloaded objects' refcnt
// never decreases, so DRAM isn't fully freed beyond what direct eviction
// allows. We verify eviction doesn't deadlock (can fill to capacity).
int success_puts = FillSegmentUntilEviction(
*service, ctx.client_id, "evict_c_", object_size, 1024 * 16 + 50);
EXPECT_GE(success_puts, 1024 * 16)
<< "Combo C: eviction should work with force-evict enabled";
std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
service->RemoveAll();
}
// =============================================================================
// Combo D: offload_force_evict=true only (should be no-op without on_evict)
// =============================================================================
TEST_F(OffloadOnEvictTest, ComboD_ForceEvictAloneIsIgnored) {
MasterServiceConfig config;
config.enable_offload = true;
config.offload_force_evict = true; // on_evict is false → force is ignored
config.default_kv_lease_ttl = 2000;
auto service = std::make_unique<MasterService>(config);
constexpr size_t seg_size = 1024 * 1024 * 16;
auto ctx =
PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size);
auto mount_ld = service->MountLocalDiskSegment(ctx.client_id, true);
ASSERT_TRUE(mount_ld.has_value());
// Should behave like Combo A (default: offload at PutEnd)
PutObject(*service, ctx.client_id, "key_d1");
PutObject(*service, ctx.client_id, "key_d2");
auto queued = DrainOffloadQueue(*service, ctx.client_id);
EXPECT_EQ(queued.size(), 2u)
<< "Combo D: FORCE_EVICT alone should not change default behavior";
service->RemoveAll();
}
TEST_F(OffloadOnEvictTest, ComboD_EvictionWorks) {
const uint64_t kv_lease_ttl = 2000;
MasterServiceConfig config;
config.enable_offload = true;
config.offload_force_evict = true; // on_evict is false → force is ignored
config.default_kv_lease_ttl = kv_lease_ttl;
auto service = std::make_unique<MasterService>(config);
constexpr size_t seg_size = 1024 * 1024 * 16 * 15;
constexpr size_t object_size = 1024 * 15;
auto ctx =
PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size);
int success_puts = FillSegmentUntilEviction(
*service, ctx.client_id, "evict_d_", object_size, 1024 * 16 + 50);
EXPECT_GT(success_puts, 1024 * 16)
<< "Combo D: eviction should work normally";
std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl));
service->RemoveAll();
}
} // namespace mooncake::test
int main(int argc, char** argv) {
::testing::InitGoogleTest(&argc, argv);
return RUN_ALL_TESTS();
}

View File

@ -1,436 +0,0 @@
#include <glog/logging.h>
#include <gtest/gtest.h>
#include <chrono>
#include <thread>
#include <vector>
#include "client_metric.h"
namespace mooncake::test {
class SsdMetricsTest : public ::testing::Test {
protected:
void SetUp() override {
google::InitGoogleLogging("SsdMetricsTest");
FLAGS_logtostderr = true;
}
void TearDown() override { google::ShutdownGoogleLogging(); }
};
TEST_F(SsdMetricsTest, InitialValuesTest) {
SsdMetric metrics;
ASSERT_EQ(metrics.ssd_read_bytes.value(), 0);
ASSERT_EQ(metrics.ssd_write_bytes.value(), 0);
ASSERT_EQ(metrics.ssd_read_ops.value(), 0);
ASSERT_EQ(metrics.ssd_write_ops.value(), 0);
// Histogram bucket counts should all be 0
auto read_buckets = metrics.ssd_read_latency_us.get_bucket_counts();
for (auto& bucket : read_buckets) {
ASSERT_EQ(bucket->value(), 0);
}
auto write_buckets = metrics.ssd_write_latency_us.get_bucket_counts();
for (auto& bucket : write_buckets) {
ASSERT_EQ(bucket->value(), 0);
}
// Total metrics
ASSERT_EQ(metrics.ssd_total_bytes.value(), 0);
ASSERT_EQ(metrics.ssd_total_ops.value(), 0);
}
TEST_F(SsdMetricsTest, ReadMetricsTest) {
SsdMetric metrics;
// Simulate a successful BatchLoad of 3 keys totaling 1MB
metrics.ssd_read_ops.inc(3);
metrics.ssd_read_bytes.inc(1024 * 1024);
metrics.ssd_read_latency_us.observe(1500); // 1.5ms
ASSERT_EQ(metrics.ssd_read_ops.value(), 3);
ASSERT_EQ(metrics.ssd_read_bytes.value(), 1024 * 1024);
// Simulate another BatchLoad of 5 keys totaling 2MB
metrics.ssd_read_ops.inc(5);
metrics.ssd_read_bytes.inc(2 * 1024 * 1024);
metrics.ssd_read_latency_us.observe(3000); // 3ms
ASSERT_EQ(metrics.ssd_read_ops.value(), 8);
ASSERT_EQ(metrics.ssd_read_bytes.value(), 3 * 1024 * 1024);
// Write metrics should remain 0
ASSERT_EQ(metrics.ssd_write_ops.value(), 0);
ASSERT_EQ(metrics.ssd_write_bytes.value(), 0);
}
TEST_F(SsdMetricsTest, WriteMetricsTest) {
SsdMetric metrics;
// Simulate a successful BatchOffload of 10 keys totaling 5MB
metrics.ssd_write_ops.inc(10);
metrics.ssd_write_bytes.inc(5 * 1024 * 1024);
metrics.ssd_write_latency_us.observe(50000); // 50ms
ASSERT_EQ(metrics.ssd_write_ops.value(), 10);
ASSERT_EQ(metrics.ssd_write_bytes.value(), 5 * 1024 * 1024);
// Read metrics should remain 0
ASSERT_EQ(metrics.ssd_read_ops.value(), 0);
ASSERT_EQ(metrics.ssd_read_bytes.value(), 0);
}
TEST_F(SsdMetricsTest, TotalMetricsTest) {
SsdMetric metrics;
// Simulate read: 3 keys, 1MB
metrics.ssd_read_ops.inc(3);
metrics.ssd_read_bytes.inc(1024 * 1024);
metrics.ssd_read_latency_us.observe(1500);
metrics.ssd_read_latency_summary.observe(1500);
metrics.ssd_total_ops.inc(3);
metrics.ssd_total_bytes.inc(1024 * 1024);
metrics.ssd_total_latency_us.observe(1500);
metrics.ssd_total_latency_summary.observe(1500);
// Simulate write: 5 keys, 2MB
metrics.ssd_write_ops.inc(5);
metrics.ssd_write_bytes.inc(2 * 1024 * 1024);
metrics.ssd_write_latency_us.observe(3000);
metrics.ssd_write_latency_summary.observe(3000);
metrics.ssd_total_ops.inc(5);
metrics.ssd_total_bytes.inc(2 * 1024 * 1024);
metrics.ssd_total_latency_us.observe(3000);
metrics.ssd_total_latency_summary.observe(3000);
// Verify totals = read + write
ASSERT_EQ(metrics.ssd_total_ops.value(), 3 + 5);
ASSERT_EQ(metrics.ssd_total_bytes.value(), 1024 * 1024 + 2 * 1024 * 1024);
// Total latency histogram should have 2 observations
auto buckets = metrics.ssd_total_latency_us.get_bucket_counts();
int64_t total_count = 0;
for (auto& b : buckets) total_count += b->value();
ASSERT_EQ(total_count, 2);
// Verify summary shows Total line
std::string summary = metrics.summary_metrics();
EXPECT_TRUE(summary.find("SSD Total:") != std::string::npos);
EXPECT_TRUE(summary.find("Total:") != std::string::npos);
std::cout << "Total Metrics Summary:\n" << summary << std::endl;
}
TEST_F(SsdMetricsTest, FailureNotCountedTest) {
SsdMetric metrics;
// Simulate: nothing recorded (as if BatchLoad failed and we skipped
// metrics) All metrics should remain 0
ASSERT_EQ(metrics.ssd_read_ops.value(), 0);
ASSERT_EQ(metrics.ssd_read_bytes.value(), 0);
ASSERT_EQ(metrics.ssd_write_ops.value(), 0);
ASSERT_EQ(metrics.ssd_write_bytes.value(), 0);
// Now simulate a success after a failure
metrics.ssd_read_ops.inc(1);
metrics.ssd_read_bytes.inc(4096);
metrics.ssd_read_latency_us.observe(100);
// Only the successful operation should be counted
ASSERT_EQ(metrics.ssd_read_ops.value(), 1);
ASSERT_EQ(metrics.ssd_read_bytes.value(), 4096);
}
TEST_F(SsdMetricsTest, ConcurrentTest) {
SsdMetric metrics;
const int num_threads = 8;
const int ops_per_thread = 1000;
std::vector<std::thread> threads;
for (int i = 0; i < num_threads; ++i) {
threads.emplace_back([&metrics, ops_per_thread]() {
for (int j = 0; j < ops_per_thread; ++j) {
metrics.ssd_read_ops.inc(1);
metrics.ssd_read_bytes.inc(4096);
metrics.ssd_read_latency_us.observe(500);
metrics.ssd_write_ops.inc(1);
metrics.ssd_write_bytes.inc(8192);
metrics.ssd_write_latency_us.observe(1000);
metrics.ssd_total_ops.inc(2);
metrics.ssd_total_bytes.inc(4096 + 8192);
}
});
}
for (auto& t : threads) {
t.join();
}
int64_t expected_ops = num_threads * ops_per_thread;
ASSERT_EQ(metrics.ssd_read_ops.value(), expected_ops);
ASSERT_EQ(metrics.ssd_read_bytes.value(), expected_ops * 4096);
ASSERT_EQ(metrics.ssd_write_ops.value(), expected_ops);
ASSERT_EQ(metrics.ssd_write_bytes.value(), expected_ops * 8192);
ASSERT_EQ(metrics.ssd_total_ops.value(), expected_ops * 2);
ASSERT_EQ(metrics.ssd_total_bytes.value(), expected_ops * (4096 + 8192));
}
TEST_F(SsdMetricsTest, SerializeTest) {
SsdMetric metrics;
// Add some data so serialization produces output
metrics.ssd_read_ops.inc(5);
metrics.ssd_read_bytes.inc(1024);
metrics.ssd_read_latency_us.observe(200);
metrics.ssd_write_ops.inc(3);
metrics.ssd_write_bytes.inc(2048);
metrics.ssd_write_latency_us.observe(1000);
metrics.ssd_total_ops.inc(8);
metrics.ssd_total_bytes.inc(3072);
metrics.ssd_total_latency_us.observe(200);
metrics.ssd_total_latency_us.observe(1000);
metrics.ssd_read_latency_summary.observe(200);
metrics.ssd_write_latency_summary.observe(1000);
metrics.ssd_total_latency_summary.observe(200);
metrics.ssd_total_latency_summary.observe(1000);
std::string serialized;
metrics.serialize(serialized);
// Verify all metric names appear in serialized output
EXPECT_TRUE(serialized.find("mooncake_ssd_read_bytes_total") !=
std::string::npos);
EXPECT_TRUE(serialized.find("mooncake_ssd_write_bytes_total") !=
std::string::npos);
EXPECT_TRUE(serialized.find("mooncake_ssd_read_ops_total") !=
std::string::npos);
EXPECT_TRUE(serialized.find("mooncake_ssd_write_ops_total") !=
std::string::npos);
EXPECT_TRUE(serialized.find("mooncake_ssd_read_latency_us") !=
std::string::npos);
EXPECT_TRUE(serialized.find("mooncake_ssd_write_latency_us") !=
std::string::npos);
EXPECT_TRUE(serialized.find("mooncake_ssd_total_bytes_total") !=
std::string::npos);
EXPECT_TRUE(serialized.find("mooncake_ssd_total_ops_total") !=
std::string::npos);
EXPECT_TRUE(serialized.find("mooncake_ssd_total_latency_us") !=
std::string::npos);
EXPECT_TRUE(serialized.find("mooncake_ssd_read_latency_summary_us") !=
std::string::npos);
std::cout << "Serialized SSD Metrics:\n" << serialized << std::endl;
}
TEST_F(SsdMetricsTest, SummaryMetricsTest) {
SsdMetric metrics;
// Test empty metrics - no throughput shown when bytes=0
std::string summary = metrics.summary_metrics();
EXPECT_TRUE(summary.find("SSD Metrics Summary") != std::string::npos);
EXPECT_TRUE(summary.find("SSD Read: 0 B") != std::string::npos);
EXPECT_TRUE(summary.find("SSD Write: 0 B") != std::string::npos);
EXPECT_TRUE(summary.find("Read: No data") != std::string::npos);
EXPECT_TRUE(summary.find("Write: No data") != std::string::npos);
// No throughput/IOPS when no data
EXPECT_TRUE(summary.find("throughput=") == std::string::npos);
EXPECT_TRUE(summary.find("IOPS=") == std::string::npos);
// Add some data
metrics.ssd_read_bytes.inc(5 * 1024 * 1024); // 5MB
metrics.ssd_write_bytes.inc(10 * 1024 * 1024); // 10MB
metrics.ssd_read_ops.inc(100);
metrics.ssd_write_ops.inc(50);
// Add latency observations (both histogram and summary)
for (int i = 0; i < 50; ++i) {
metrics.ssd_read_latency_us.observe(500);
metrics.ssd_read_latency_summary.observe(500);
}
for (int i = 0; i < 40; ++i) {
metrics.ssd_read_latency_us.observe(2000);
metrics.ssd_read_latency_summary.observe(2000);
}
for (int i = 0; i < 10; ++i) {
metrics.ssd_read_latency_us.observe(50000);
metrics.ssd_read_latency_summary.observe(50000);
}
summary = metrics.summary_metrics();
EXPECT_TRUE(summary.find("SSD Read: 5.00 MB") != std::string::npos);
EXPECT_TRUE(summary.find("SSD Write: 10.00 MB") != std::string::npos);
EXPECT_TRUE(summary.find("ops=100") != std::string::npos);
EXPECT_TRUE(summary.find("ops=50") != std::string::npos);
// Throughput and IOPS should now appear
EXPECT_TRUE(summary.find("throughput=") != std::string::npos);
EXPECT_TRUE(summary.find("/s") != std::string::npos);
EXPECT_TRUE(summary.find("IOPS=") != std::string::npos);
// Latency percentiles
EXPECT_TRUE(summary.find("p50=") != std::string::npos);
EXPECT_TRUE(summary.find("p90=") != std::string::npos);
EXPECT_TRUE(summary.find("p99=") != std::string::npos);
std::cout << "SSD Metrics Summary:\n" << summary << std::endl;
}
TEST_F(SsdMetricsTest, ThroughputCalculationTest) {
SsdMetric metrics;
// Record 10MB of reads and 20MB of writes
int64_t read_bytes = 10 * 1024 * 1024;
int64_t write_bytes = 20 * 1024 * 1024;
metrics.ssd_read_bytes.inc(read_bytes);
metrics.ssd_write_bytes.inc(write_bytes);
metrics.ssd_read_ops.inc(1000);
metrics.ssd_write_ops.inc(2000);
metrics.ssd_read_latency_us.observe(1000);
metrics.ssd_read_latency_summary.observe(1000);
metrics.ssd_write_latency_us.observe(2000);
metrics.ssd_write_latency_summary.observe(2000);
metrics.ssd_total_bytes.inc(read_bytes + write_bytes);
metrics.ssd_total_ops.inc(3000);
metrics.ssd_total_latency_us.observe(1000);
metrics.ssd_total_latency_us.observe(2000);
metrics.ssd_total_latency_summary.observe(1000);
metrics.ssd_total_latency_summary.observe(2000);
// Wait a small amount to get non-zero elapsed time
std::this_thread::sleep_for(std::chrono::milliseconds(100));
std::string summary = metrics.summary_metrics();
// Verify throughput is displayed for both read and write
// Read line should have: throughput=<value>/s, IOPS=<value>
EXPECT_TRUE(summary.find("throughput=") != std::string::npos)
<< "Missing throughput in summary:\n"
<< summary;
EXPECT_TRUE(summary.find("IOPS=") != std::string::npos)
<< "Missing IOPS in summary:\n"
<< summary;
// Count occurrences of "throughput=" - should be 3 (read + write + total)
size_t count = 0;
size_t pos = 0;
while ((pos = summary.find("throughput=", pos)) != std::string::npos) {
++count;
pos += 11;
}
EXPECT_EQ(count, 3) << "Expected throughput for read, write, and total";
// Count occurrences of "IOPS=" - should be 3 (read + write + total)
count = 0;
pos = 0;
while ((pos = summary.find("IOPS=", pos)) != std::string::npos) {
++count;
pos += 5;
}
EXPECT_EQ(count, 3) << "Expected IOPS for read, write, and total";
std::cout << "Throughput Test Summary:\n" << summary << std::endl;
}
TEST_F(SsdMetricsTest, IntegrationWithClientMetric) {
ClientMetric client_metrics;
// Add SSD data via client_metrics.ssd_metric
client_metrics.ssd_metric.ssd_read_ops.inc(10);
client_metrics.ssd_metric.ssd_read_bytes.inc(1024 * 1024);
client_metrics.ssd_metric.ssd_read_latency_us.observe(1500);
// Verify it appears in ClientMetric serialize
std::string serialized;
client_metrics.serialize(serialized);
EXPECT_TRUE(serialized.find("mooncake_ssd_read_ops_total") !=
std::string::npos);
// Verify it appears in ClientMetric summary
std::string summary = client_metrics.summary_metrics();
EXPECT_TRUE(summary.find("SSD Metrics Summary") != std::string::npos);
std::cout << "Client Metrics with SSD:\n" << summary << std::endl;
}
TEST_F(SsdMetricsTest, SerializeWithDynamicLabels) {
std::map<std::string, std::string> labels = {{"instance_id", "test123"},
{"cluster_id", "cluster_abc"}};
SsdMetric metrics(labels);
metrics.ssd_read_ops.inc(1);
metrics.ssd_read_latency_us.observe(100);
std::string serialized;
metrics.serialize(serialized);
EXPECT_TRUE(serialized.find("instance_id=\"test123\"") !=
std::string::npos);
EXPECT_TRUE(serialized.find("cluster_id=\"cluster_abc\"") !=
std::string::npos);
std::cout << "SSD Metrics with labels:\n" << serialized << std::endl;
}
TEST_F(SsdMetricsTest, LatencyBucketBoundaryTest) {
SsdMetric metrics;
// kSsdLatencyBucket = {50, 100, 200, 500, 1000, ... 10000000, 30000000}
// Value below the smallest bucket (50us)
metrics.ssd_read_latency_us.observe(25);
// Value exactly at bucket boundary (50us)
metrics.ssd_read_latency_us.observe(50);
// Value at the largest bucket boundary (30s)
metrics.ssd_read_latency_us.observe(30000000);
// Value exceeding all buckets -> falls into +Inf
metrics.ssd_read_latency_us.observe(60000000);
auto buckets = metrics.ssd_read_latency_us.get_bucket_counts();
ASSERT_FALSE(buckets.empty());
// First bucket (le=50): should contain 2 observations (25 and 50)
EXPECT_EQ(buckets[0]->value(), 2);
// Last real bucket (le=30000000): should contain 1 observation (30000000)
// (buckets[0..N-2] are real buckets, buckets[N-1] is +Inf)
size_t last_real_idx = kSsdLatencyBucket.size() - 1;
EXPECT_EQ(buckets[last_real_idx]->value(), 1);
// +Inf bucket: should contain 1 observation (60000000)
size_t inf_idx = buckets.size() - 1;
EXPECT_EQ(buckets[inf_idx]->value(), 1);
// Verify total count via serialization
std::string serialized;
metrics.serialize(serialized);
EXPECT_TRUE(serialized.find("mooncake_ssd_read_latency_us_count 4") !=
std::string::npos);
std::cout << "Bucket boundary test serialization:\n"
<< serialized << std::endl;
}
TEST_F(SsdMetricsTest, EmptyBatchMetricsTest) {
SsdMetric metrics;
// Simulate recording metrics for an empty batch (0 keys, 0 bytes)
// This mirrors what happens if FileStorage::BatchLoad is called with
// an empty map and succeeds - the instrumentation code would compute
// total_bytes=0 and batch_object.size()=0
metrics.ssd_read_ops.inc(0);
metrics.ssd_read_bytes.inc(0);
metrics.ssd_read_latency_us.observe(10); // Still takes some time
// Counters should remain 0 (inc(0) is a no-op for counters)
ASSERT_EQ(metrics.ssd_read_ops.value(), 0);
ASSERT_EQ(metrics.ssd_read_bytes.value(), 0);
// But latency histogram should record the observation
auto buckets = metrics.ssd_read_latency_us.get_bucket_counts();
int64_t total = 0;
for (auto& b : buckets) total += b->value();
EXPECT_EQ(total, 1);
}
} // namespace mooncake::test

View File

@ -9,11 +9,10 @@ if (NOT GLOBAL_CONFIG)
include(../mooncake-common/FindGLOG.cmake)
include(../mooncake-common/common.cmake)
# Build asio_shared in mooncake-common for standalone builds.
# Set compile definitions before adding targets that include ASIO headers.
# Build asio_shared library for standalone build
# Set compile definitions before add_subdirectory to ensure proper compilation
add_compile_definitions(ASIO_SEPARATE_COMPILATION ASIO_DYN_LINK)
add_subdirectory(../mooncake-common/src
${CMAKE_CURRENT_BINARY_DIR}/mooncake-common-src)
add_subdirectory(../mooncake-asio ${CMAKE_CURRENT_BINARY_DIR}/mooncake-asio)
endif() # GLOBAL_CONFIG
if (USE_ASCEND)

View File

@ -9,12 +9,11 @@ endif()
file(GLOB TEBENCH_SOURCES "*.cpp")
add_executable(tebench ${TEBENCH_SOURCES})
target_link_libraries(tebench PUBLIC transfer_engine tent_link_group)
target_link_libraries(tebench PUBLIC transfer_engine tent)
if(USE_CUDA)
target_link_libraries(tebench PUBLIC CUDA::cudart)
endif()
# Set RPATH for finding libasio.so at runtime
set_target_properties(
tebench PROPERTIES BUILD_WITH_INSTALL_RPATH TRUE
INSTALL_RPATH "$ORIGIN/../lib:$ORIGIN/../../mooncake-common")
set_target_properties(tebench PROPERTIES BUILD_WITH_INSTALL_RPATH TRUE
INSTALL_RPATH "$ORIGIN/../lib:$ORIGIN/../../mooncake-asio")

View File

@ -131,12 +131,7 @@ int TENTBenchRunner::allocateBuffers() {
// Register
auto allocated_ts = getCurrentTimeInNano();
std::vector<size_t> buffers_size(num_buffers, total_buffer_size);
MemoryOptions options;
if (!xport_type.empty()) {
options.type = getTransportType(xport_type);
}
CHECK_FAIL(engine_->registerLocalMemory(pinned_buffer_list_, buffers_size,
options));
CHECK_FAIL(engine_->registerLocalMemory(pinned_buffer_list_, buffers_size));
auto registered_ts = getCurrentTimeInNano();
LOG(INFO) << "Allocated " << total_buffer_size * num_buffers << " bytes "

View File

@ -1,57 +1,47 @@
set(WORKSPACE "${CMAKE_CURRENT_SOURCE_DIR}")
if(USE_HIP)
if (USE_HIP)
file(GLOB EXAMPLE_SOURCES "*.cpp")
hipify_files(EXAMPLE_SOURCES)
file(RELATIVE_PATH EXAMPLE_REL_PATH "${CMAKE_SOURCE_DIR}"
"${CMAKE_CURRENT_SOURCE_DIR}")
file(RELATIVE_PATH EXAMPLE_REL_PATH "${CMAKE_SOURCE_DIR}" "${CMAKE_CURRENT_SOURCE_DIR}")
set(WORKSPACE "${CMAKE_BINARY_DIR}/${EXAMPLE_REL_PATH}")
endif()
add_executable(transfer_engine_bench ${WORKSPACE}/transfer_engine_bench.cpp)
target_link_libraries(transfer_engine_bench PUBLIC transfer_engine)
if(USE_TENT)
target_link_libraries(transfer_engine_bench PUBLIC tent_link_group)
target_compile_definitions(transfer_engine_bench PRIVATE USE_TENT)
if (USE_TENT)
target_link_libraries(transfer_engine_bench PUBLIC tent)
target_compile_definitions(transfer_engine_bench PRIVATE USE_TENT)
endif()
add_executable(transfer_engine_validator
${WORKSPACE}/transfer_engine_validator.cpp)
add_executable(transfer_engine_validator ${WORKSPACE}/transfer_engine_validator.cpp)
target_link_libraries(transfer_engine_validator PUBLIC transfer_engine)
add_executable(transfer_engine_bench_with_notify
${WORKSPACE}/transfer_engine_bench_with_notify.cpp)
add_executable(transfer_engine_bench_with_notify ${WORKSPACE}/transfer_engine_bench_with_notify.cpp)
target_link_libraries(transfer_engine_bench_with_notify PUBLIC transfer_engine)
add_executable(memory_pool ${WORKSPACE}/memory_pool.cpp)
target_link_libraries(memory_pool PUBLIC transfer_engine)
if(USE_ASCEND)
add_executable(transfer_engine_ascend_one_sided
${WORKSPACE}/transfer_engine_ascend_one_sided.cpp)
target_link_libraries(transfer_engine_ascend_one_sided PUBLIC transfer_engine)
if (USE_ASCEND)
add_executable(transfer_engine_ascend_one_sided ${WORKSPACE}/transfer_engine_ascend_one_sided.cpp)
target_link_libraries(transfer_engine_ascend_one_sided PUBLIC transfer_engine)
add_executable(transfer_engine_ascend_perf
${WORKSPACE}/transfer_engine_ascend_perf.cpp)
target_link_libraries(transfer_engine_ascend_perf PUBLIC transfer_engine)
add_executable(transfer_engine_ascend_perf ${WORKSPACE}/transfer_engine_ascend_perf.cpp)
target_link_libraries(transfer_engine_ascend_perf PUBLIC transfer_engine)
endif()
if(USE_ASCEND_DIRECT)
add_executable(transfer_engine_ascend_direct_perf
${WORKSPACE}/transfer_engine_ascend_direct_perf.cpp)
target_link_libraries(transfer_engine_ascend_direct_perf
PUBLIC ascendcl transfer_engine)
if (USE_ASCEND_DIRECT)
add_executable(transfer_engine_ascend_direct_perf ${WORKSPACE}/transfer_engine_ascend_direct_perf.cpp)
target_link_libraries(transfer_engine_ascend_direct_perf PUBLIC ascendcl transfer_engine)
endif()
if(USE_ASCEND_HETEROGENEOUS)
add_executable(
transfer_engine_heterogeneous_ascend_perf_initiator
${WORKSPACE}/transfer_engine_heterogeneous_ascend_perf_initiator.cpp)
target_link_libraries(transfer_engine_heterogeneous_ascend_perf_initiator
PUBLIC transfer_engine)
if (USE_ASCEND_HETEROGENEOUS)
add_executable(transfer_engine_heterogeneous_ascend_perf_initiator ${WORKSPACE}/transfer_engine_heterogeneous_ascend_perf_initiator.cpp)
target_link_libraries(transfer_engine_heterogeneous_ascend_perf_initiator PUBLIC transfer_engine)
endif()
if(USE_UBSHMEM)
target_link_libraries(transfer_engine_bench PUBLIC transfer_engine)
endif()
if (USE_UBSHMEM)
target_link_libraries(transfer_engine_bench PUBLIC transfer_engine)
endif()

View File

@ -1,357 +0,0 @@
#!/usr/bin/env python3
"""
Batch Memory Registration Benchmark for Mooncake EFA Transport.
Tests registering multiple independent memory blocks (simulating
multi-tenant or sharded KV cache pools) and measures registration
time and transfer throughput across blocks.
Usage:
# Target node (registers memory blocks):
python batch_register_bench.py --mode target \
--local_server_name <target_ip>:12345 \
--num_blocks 40 --block_size_gb 4 --protocol efa
# Initiator node (pulls data from random blocks):
python batch_register_bench.py --mode initiator \
--local_server_name <initiator_ip>:12346 \
--target_server_name <target_ip>:12345 \
--num_blocks 40 --block_size_gb 4 --protocol efa
Requires: mooncake Python package (pip install -e mooncake-wheel)
"""
import argparse
import ctypes
import ctypes.util
import json
import os
import random
import signal
import statistics
import sys
import time
def parse_args():
parser = argparse.ArgumentParser(
description="Batch Memory Registration Benchmark"
)
parser.add_argument(
"--mode",
choices=["target", "initiator"],
required=True,
help="Run as target (memory holder) or initiator (puller)",
)
parser.add_argument(
"--local_server_name",
required=True,
help="Local address, e.g. 172.31.6.162:12345",
)
parser.add_argument(
"--target_server_name",
default="",
help="Target address (initiator mode only)",
)
parser.add_argument(
"--metadata_server",
default="P2PHANDSHAKE",
help="Metadata server address (default: P2PHANDSHAKE)",
)
parser.add_argument(
"--protocol", default="efa", help="Transport protocol (default: efa)"
)
parser.add_argument(
"--num_blocks",
type=int,
default=40,
help="Number of memory blocks to register (default: 40)",
)
parser.add_argument(
"--block_size_gb",
type=float,
default=4.0,
help="Size of each memory block in GB (default: 4)",
)
parser.add_argument(
"--transfer_size_mb",
type=float,
default=368.0,
help="Transfer size in MB per iteration (default: 368, ~4K tokens GLM-5.1)",
)
parser.add_argument(
"--iterations",
type=int,
default=20,
help="Number of transfer iterations (default: 20)",
)
parser.add_argument(
"--warmup",
type=int,
default=5,
help="Number of warmup iterations (default: 5)",
)
parser.add_argument(
"--use_batch_api",
action="store_true",
help="Use batch_register_memory API instead of per-block register_memory",
)
return parser.parse_args()
def allocate_block(size_bytes):
"""Allocate a single page-aligned memory block using mmap."""
libc_name = ctypes.util.find_library("c")
libc = ctypes.CDLL(libc_name, use_errno=True)
PROT_READ = 0x1
PROT_WRITE = 0x2
MAP_PRIVATE = 0x02
MAP_ANONYMOUS = 0x20
MAP_HUGETLB = 0x40000
MAP_FAILED = ctypes.c_void_p(-1).value
libc.mmap.restype = ctypes.c_void_p
libc.mmap.argtypes = [
ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int,
ctypes.c_int, ctypes.c_int, ctypes.c_long,
]
flags = MAP_PRIVATE | MAP_ANONYMOUS
# Try hugepages first
ptr = libc.mmap(None, size_bytes, PROT_READ | PROT_WRITE,
flags | MAP_HUGETLB, -1, 0)
if ptr and ptr != MAP_FAILED:
return ptr, True
# Fall back to regular pages
ptr = libc.mmap(None, size_bytes, PROT_READ | PROT_WRITE, flags, -1, 0)
if not ptr or ptr == MAP_FAILED:
raise RuntimeError(
f"mmap failed for {size_bytes} bytes: errno={ctypes.get_errno()}"
)
return ptr, False
def run_target(args):
"""Run as target: allocate blocks, register, and wait."""
from mooncake.engine import TransferEngine
block_bytes = int(args.block_size_gb * 1024 * 1024 * 1024)
total_gb = args.num_blocks * args.block_size_gb
print(f"=== Target Node ===")
print(f"Blocks: {args.num_blocks} x {args.block_size_gb} GB = {total_gb} GB total")
print(f"Protocol: {args.protocol}")
print(f"Registration API: {'batch' if args.use_batch_api else 'per-block'}")
engine = TransferEngine()
ret = engine.initialize(
args.local_server_name, args.metadata_server, args.protocol, ""
)
if ret != 0:
raise RuntimeError(f"Engine initialization failed: {ret}")
if args.metadata_server == "P2PHANDSHAKE":
host = args.local_server_name.rpartition(":")[0]
rpc_port = engine.get_rpc_port()
actual_name = f"{host}:{rpc_port}"
print(f"Actual server name: {actual_name}")
# Allocate blocks
print(f"\nAllocating {args.num_blocks} blocks of {args.block_size_gb} GB...")
blocks = [] # (addr, size, is_hugepage)
for i in range(args.num_blocks):
try:
ptr, hp = allocate_block(block_bytes)
blocks.append((ptr, block_bytes, hp))
if (i + 1) % 10 == 0 or i == 0:
page_type = "hugepage" if hp else "4KB"
print(f" Block {i+1}/{args.num_blocks}: 0x{ptr:x} ({page_type})")
except RuntimeError as e:
print(f" Block {i+1} allocation FAILED: {e}")
break
if not blocks:
raise RuntimeError("No blocks allocated")
hp_count = sum(1 for _, _, hp in blocks if hp)
reg_count = sum(1 for _, _, hp in blocks if not hp)
print(f" Allocated {len(blocks)} blocks: {hp_count} hugepage, {reg_count} regular")
# Register
addrs = [b[0] for b in blocks]
sizes = [b[1] for b in blocks]
print(f"\nRegistering {len(blocks)} blocks...")
t0 = time.time()
if args.use_batch_api:
ret = engine.batch_register_memory(addrs, sizes)
if ret != 0:
print(f" batch_register_memory FAILED: {ret}")
print(f" Registration took {time.time() - t0:.1f}s before failure")
return
else:
for i, (addr, size, _) in enumerate(blocks):
ret = engine.register_memory(addr, size)
if ret != 0:
print(f" register_memory FAILED at block {i}: {ret}")
print(f" Registration took {time.time() - t0:.1f}s before failure")
return
reg_time = time.time() - t0
print(f" Registration OK: {reg_time:.1f}s for {len(blocks)} blocks")
print(f" Per-block: {reg_time / len(blocks) * 1000:.0f}ms")
# Publish block info
info = {
"blocks": [{"addr": b[0], "size": b[1]} for b in blocks],
"num_blocks": len(blocks),
}
print(f"\nTarget ready with {len(blocks)} blocks.")
print(f"TARGET_INFO:{json.dumps(info)}")
print("Waiting for initiator (Ctrl+C to stop)...")
try:
signal.pause()
except KeyboardInterrupt:
print("\nShutting down target.")
def run_initiator(args):
"""Run as initiator: pull data from random blocks."""
from mooncake.engine import TransferEngine
block_bytes = int(args.block_size_gb * 1024 * 1024 * 1024)
transfer_bytes = int(args.transfer_size_mb * 1024 * 1024)
print(f"=== Initiator Node ===")
print(f"Target: {args.target_server_name}")
print(f"Blocks: {args.num_blocks} x {args.block_size_gb} GB")
print(f"Transfer size: {args.transfer_size_mb} MB")
print(f"Protocol: {args.protocol}")
if not args.target_server_name:
raise RuntimeError("--target_server_name required in initiator mode")
engine = TransferEngine()
ret = engine.initialize(
args.local_server_name, args.metadata_server, args.protocol, ""
)
if ret != 0:
raise RuntimeError(f"Engine initialization failed: {ret}")
if args.metadata_server == "P2PHANDSHAKE":
host = args.local_server_name.rpartition(":")[0]
rpc_port = engine.get_rpc_port()
actual_name = f"{host}:{rpc_port}"
print(f"Actual server name: {actual_name}")
# Allocate local receive buffer
recv_bytes = transfer_bytes
print(f"\nAllocating {recv_bytes / 1e6:.0f} MB receive buffer...")
recv_addr, hp = allocate_block(recv_bytes)
print(f" 0x{recv_addr:x} ({'hugepage' if hp else '4KB'})")
ret = engine.register_memory(recv_addr, recv_bytes)
if ret != 0:
raise RuntimeError(f"Local memory registration failed: {ret}")
# Get target's first buffer address
print(f"Connecting to target {args.target_server_name}...")
remote_base = engine.get_first_buffer_address(args.target_server_name)
if remote_base == 0:
raise RuntimeError("Cannot get target buffer address")
print(f" Remote first buffer at 0x{remote_base:x}")
# Connection warmup
print("Warming up connection...")
for _ in range(3):
engine.transfer_sync_read(
args.target_server_name, recv_addr, remote_base, transfer_bytes
)
print(" Connection ready.")
# Note: blocks are independently mmap'd on the target, so they are NOT
# contiguous. We can only access the first block via get_first_buffer_address.
# The primary goal is to validate that registration of N separate blocks works
# and that data can be transferred from a registered block.
print(f"\nBenchmarking transfers from first registered block...")
print(f" Each transfer: {args.transfer_size_mb} MB")
# Warmup
print(f" Warming up ({args.warmup} iterations)...")
for w in range(args.warmup):
ret = engine.transfer_sync_read(
args.target_server_name, recv_addr, remote_base, transfer_bytes
)
if ret != 0:
print(f" WARNING: warmup failed iter {w}: {ret}")
# Benchmark
latencies = []
errors = 0
for i in range(args.iterations):
t0 = time.perf_counter()
ret = engine.transfer_sync_read(
args.target_server_name, recv_addr, remote_base, transfer_bytes
)
elapsed = time.perf_counter() - t0
if ret != 0:
errors += 1
if errors <= 3:
print(f" ERROR: transfer failed iter {i}: {ret}")
continue
latencies.append(elapsed * 1000)
if not latencies:
print(f" ALL FAILED ({errors} errors)")
return
latencies.sort()
avg_ms = statistics.mean(latencies)
p50_ms = latencies[len(latencies) // 2]
p99_ms = latencies[int(len(latencies) * 0.99)]
throughput_gbs = (transfer_bytes / 1e9) / (p50_ms / 1000)
print(f"\n{'='*60}")
print(f" Registered blocks: {args.num_blocks} x {args.block_size_gb}GB")
print(f" Transfer: {args.transfer_size_mb} MB")
print(f" Iterations: {len(latencies)} (errors: {errors})")
print(f" Avg latency: {avg_ms:.2f} ms")
print(f" p50 latency: {p50_ms:.2f} ms")
print(f" p99 latency: {p99_ms:.2f} ms")
print(f" Throughput: {throughput_gbs:.2f} GB/s")
print(f"{'='*60}")
results = {
"num_blocks": args.num_blocks,
"block_size_gb": args.block_size_gb,
"transfer_size_mb": args.transfer_size_mb,
"iterations": len(latencies),
"errors": errors,
"avg_latency_ms": round(avg_ms, 3),
"p50_latency_ms": round(p50_ms, 3),
"p99_latency_ms": round(p99_ms, 3),
"throughput_gbs": round(throughput_gbs, 3),
}
out_file = f"batch_bench_{args.num_blocks}x{args.block_size_gb}gb.json"
with open(out_file, "w") as f:
json.dump(results, f, indent=2)
print(f"Results saved to {out_file}")
def main():
args = parse_args()
if args.mode == "target":
run_target(args)
else:
run_initiator(args)
if __name__ == "__main__":
main()

View File

@ -1,241 +0,0 @@
#!/usr/bin/env python3
"""
EFA Per-Transfer Latency Benchmark
Measures single-transfer latency by running transfer_engine_bench with
threads=1, batch_size=1 across multiple block sizes. This isolates the
per-request overhead (slice creation, spinlock, atomic ops, MR lookup)
that the P0 NIC-striping optimization aims to eliminate.
Usage:
python3 efa_per_transfer_latency_bench.py \
--target_host=HOST_A --initiator_host=HOST_B \
--build_dir=/path/to/build \
--ssh_opts="-i /path/to/key.pem"
"""
import argparse
import os
import re
import subprocess
import sys
import time
def parse_args():
parser = argparse.ArgumentParser(
description="EFA Per-Transfer Latency Benchmark"
)
parser.add_argument("--target_host", required=True)
parser.add_argument("--initiator_host", required=True)
parser.add_argument(
"--build_dir",
default="/opt/dlami/nvme/Mooncake/build",
)
parser.add_argument("--duration", type=int, default=10)
parser.add_argument("--operation", default="write", choices=["read", "write"])
parser.add_argument("--ssh_user", default="ubuntu")
parser.add_argument(
"--ssh_opts",
default="-o StrictHostKeyChecking=no -o ConnectTimeout=10",
)
parser.add_argument(
"--block_sizes",
default="65536,131072,262144,524288,1048576,2097152,4194304,8388608,16777216",
help="Comma-separated block sizes in bytes",
)
parser.add_argument("--threads", type=int, default=1, help="Number of threads")
parser.add_argument("--batch_size", type=int, default=1, help="Batch size")
parser.add_argument(
"--env", action="append", default=[],
help="Environment variables to pass to remote bench (e.g. --env MC_EFA_STRIPING_THRESHOLD=67108864)",
)
parser.add_argument("--output", default=None, help="Output file for results")
return parser.parse_args()
def run_ssh(host, command, user, ssh_opts, timeout=None):
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 kill_bench(host, user, ssh_opts):
cmd = (
"ps aux | grep '[t]ransfer_engine_bench' | awk '{print $2}' "
"| xargs -r kill 2>/dev/null; sleep 1; echo done"
)
run_ssh(host, cmd, user, ssh_opts, timeout=15)
time.sleep(2)
def start_target(host, build_dir, user, ssh_opts):
bench_bin = os.path.join(
build_dir, "mooncake-transfer-engine/example/transfer_engine_bench"
)
log_file = "/tmp/efa_latency_target.log"
target_cmd = (
f"cd {build_dir} && "
f"{bench_bin} "
f"--mode=target --protocol=efa --metadata_server=P2PHANDSHAKE "
f"> {log_file} 2>&1"
)
ssh_args = ["ssh", "-n", *ssh_opts.split(), f"{user}@{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(
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)
return None
def run_single_bench(host, build_dir, target_addr, block_size,
duration, operation, user, ssh_opts,
threads=1, batch_size=1, env_vars=None):
"""Run bench with configurable threads and batch_size."""
bench_bin = os.path.join(
build_dir, "mooncake-transfer-engine/example/transfer_engine_bench"
)
env_prefix = " ".join(env_vars) + " " if env_vars else ""
bench_cmd = (
f"cd {build_dir} && "
f"{env_prefix}"
f"{bench_bin} "
f"--mode=initiator --protocol=efa --metadata_server=P2PHANDSHAKE "
f"--segment_id={target_addr} "
f"--operation={operation} "
f"--duration={duration} "
f"--threads={threads} "
f"--block_size={block_size} "
f"--batch_size={batch_size} "
f"2>&1"
)
timeout = duration + 60
rc, stdout, stderr = run_ssh(host, bench_cmd, user, ssh_opts, timeout=timeout)
combined = stdout + "\n" + stderr
# Parse throughput
match = re.search(r"throughput\s+([\d.]+)\s+GB/s", combined)
if match:
return float(match.group(1))
# Try MB/s
match = re.search(r"throughput\s+([\d.]+)\s+MB/s", combined)
if match:
return float(match.group(1)) / 1024.0
print(f" WARNING: Could not parse throughput", file=sys.stderr)
for line in combined.strip().split("\n")[-3:]:
print(f" {line}", file=sys.stderr)
return None
def format_size(size_bytes):
if size_bytes >= 1048576:
return f"{size_bytes / 1048576:.0f}MB"
return f"{size_bytes / 1024:.0f}KB"
def main():
args = parse_args()
block_sizes = [int(x) for x in args.block_sizes.split(",")]
print("=" * 70)
print("EFA Per-Transfer Latency Benchmark")
print("=" * 70)
print(f" Target: {args.target_host}")
print(f" Initiator: {args.initiator_host}")
print(f" Build dir: {args.build_dir}")
print(f" Duration: {args.duration}s per point")
print(f" Operation: {args.operation}")
print(f" Mode: threads={args.threads}, batch_size={args.batch_size}")
print(f" Block sizes: {[format_size(b) for b in block_sizes]}")
print()
# Start target
kill_bench(args.target_host, args.ssh_user, args.ssh_opts)
print("Starting target...", end="", flush=True)
target_addr = start_target(
args.target_host, args.build_dir, args.ssh_user, args.ssh_opts
)
if not target_addr:
print(" FAILED")
sys.exit(1)
print(f" ready ({target_addr})")
print()
results = []
for i, block_size in enumerate(block_sizes):
tag = format_size(block_size)
print(f" [{i+1}/{len(block_sizes)}] {tag:>6} ...", end="", flush=True)
tp = run_single_bench(
args.initiator_host, args.build_dir, target_addr,
block_size, args.duration, args.operation,
args.ssh_user, args.ssh_opts,
threads=args.threads, batch_size=args.batch_size,
env_vars=args.env,
)
if tp is None or tp == 0:
print(" FAILED")
results.append((block_size, None, None))
continue
# latency = block_size / throughput
tp_bytes = tp * 1e9 # GB/s -> bytes/s
latency_us = (block_size / tp_bytes) * 1e6 # microseconds
results.append((block_size, tp, latency_us))
print(f" {tp:7.2f} GB/s latency={latency_us:8.1f} us")
# Cleanup
kill_bench(args.target_host, args.ssh_user, args.ssh_opts)
# Summary
print()
print("=" * 70)
print("Results Summary")
print("=" * 70)
print(f"{'Block Size':>12} {'Throughput':>12} {'Latency (us)':>14}")
print("-" * 42)
for block_size, tp, lat in results:
tag = format_size(block_size)
if tp is not None:
print(f"{tag:>12} {tp:>9.2f} GB/s {lat:>11.1f} us")
else:
print(f"{tag:>12} {'N/A':>12} {'N/A':>14}")
# Write output file
if args.output:
with open(args.output, "w") as f:
f.write(f"# EFA Per-Transfer Latency Benchmark\n")
f.write(f"# Operation: {args.operation}\n")
f.write(f"# Duration: {args.duration}s per point\n")
f.write(f"# Mode: threads={args.threads}, batch_size={args.batch_size}\n")
f.write(f"#\n")
f.write(f"{'block_bytes':>12} {'block_size':>10} {'gbps':>10} {'latency_us':>12}\n")
for block_size, tp, lat in results:
tag = format_size(block_size)
if tp is not None:
f.write(f"{block_size:>12} {tag:>10} {tp:>10.2f} {lat:>12.1f}\n")
else:
f.write(f"{block_size:>12} {tag:>10} {'N/A':>10} {'N/A':>12}\n")
print(f"\nResults saved to: {args.output}")
if __name__ == "__main__":
main()

View File

@ -1,452 +0,0 @@
#!/usr/bin/env python3
"""
KV Cache Prefix Transfer Benchmark for Mooncake EFA Transport.
Simulates cross-node KV cache transfer for prefix cache hits in LLM inference.
Tests how registered memory pool size affects transfer latency/throughput
under the per-NIC partition auto-split strategy.
Usage:
# Target node (holds KV cache pool):
python kvcache_prefix_bench.py --mode target \
--local_server_name <target_ip>:12345 \
--pool_size_gb 100 --protocol efa
# Initiator node (pulls prefix KV cache):
python kvcache_prefix_bench.py --mode initiator \
--local_server_name <initiator_ip>:12346 \
--target_server_name <target_ip>:12345 \
--pool_size_gb 100 --protocol efa \
--prefix_tokens 4096,8192,16384,32768
Requires: mooncake Python package (pip install -e mooncake-wheel)
Branch: feat/efa-auto-split-mr (per-NIC partition auto-split)
"""
import argparse
import ctypes
import ctypes.util
import json
import os
import signal
import statistics
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
def parse_args():
parser = argparse.ArgumentParser(
description="KV Cache Prefix Transfer Benchmark"
)
parser.add_argument(
"--mode",
choices=["target", "initiator"],
required=True,
help="Run as target (KV cache holder) or initiator (puller)",
)
parser.add_argument(
"--local_server_name",
required=True,
help="Local address, e.g. 172.31.6.162:12345",
)
parser.add_argument(
"--target_server_name",
default="",
help="Target address (initiator mode only)",
)
parser.add_argument(
"--metadata_server",
default="P2PHANDSHAKE",
help="Metadata server address (default: P2PHANDSHAKE)",
)
parser.add_argument(
"--protocol", default="efa", help="Transport protocol (default: efa)"
)
parser.add_argument(
"--pool_size_gb",
type=float,
default=10.0,
help="KV cache pool size in GB to register (default: 10)",
)
parser.add_argument(
"--prefix_tokens",
default="4096,8192,16384,32768",
help="Comma-separated list of prefix token counts to test",
)
parser.add_argument(
"--kv_bytes_per_token",
type=int,
default=89856,
help="KV cache bytes per token. "
"Default: 89856 for GLM-5.1 (754B MoE, MLA attention): "
"(kv_lora_rank=512 + qk_rope_head_dim=64) * 2 bytes * 78 layers. "
"For GLM-4-9B (standard MHA): use 40960",
)
parser.add_argument(
"--iterations",
type=int,
default=50,
help="Number of transfer iterations per test (default: 50)",
)
parser.add_argument(
"--warmup",
type=int,
default=5,
help="Number of warmup iterations (default: 5)",
)
parser.add_argument(
"--use_gpu",
action="store_true",
help="Use GPU memory instead of CPU memory",
)
parser.add_argument(
"--gpu_id", type=int, default=0, help="GPU device ID (default: 0)"
)
parser.add_argument(
"--threads",
type=int,
default=1,
help="Number of threads for concurrent transfer (default: 1). "
"Each thread transfers a chunk of the prefix in parallel.",
)
return parser.parse_args()
def allocate_cpu_memory(size_bytes):
"""Allocate page-aligned CPU memory using mmap."""
libc_name = ctypes.util.find_library("c")
libc = ctypes.CDLL(libc_name, use_errno=True)
# mmap(NULL, size, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0)
PROT_READ = 0x1
PROT_WRITE = 0x2
MAP_PRIVATE = 0x02
MAP_ANONYMOUS = 0x20
MAP_HUGETLB = 0x40000
MAP_FAILED = ctypes.c_void_p(-1).value
libc.mmap.restype = ctypes.c_void_p
libc.mmap.argtypes = [
ctypes.c_void_p,
ctypes.c_size_t,
ctypes.c_int,
ctypes.c_int,
ctypes.c_int,
ctypes.c_long,
]
flags = MAP_PRIVATE | MAP_ANONYMOUS
# Try hugepages first (recommended for large EFA registrations)
ptr = libc.mmap(None, size_bytes, PROT_READ | PROT_WRITE,
flags | MAP_HUGETLB, -1, 0)
if ptr and ptr != MAP_FAILED:
print(f" Allocated {size_bytes / 1e9:.1f} GB with 2MB hugepages")
return ptr, size_bytes, True
# Fall back to regular pages
ptr = libc.mmap(None, size_bytes, PROT_READ | PROT_WRITE, flags, -1, 0)
if not ptr or ptr == MAP_FAILED:
raise RuntimeError(
f"mmap failed for {size_bytes} bytes: "
f"errno={ctypes.get_errno()}"
)
print(
f" Allocated {size_bytes / 1e9:.1f} GB with 4KB pages "
f"(hugepages unavailable — configure vm.nr_hugepages for better EFA performance)"
)
return ptr, size_bytes, False
def allocate_gpu_memory(size_bytes, gpu_id):
"""Allocate GPU memory using PyTorch."""
import torch
torch.cuda.set_device(gpu_id)
# Allocate as uint8 tensor
tensor = torch.zeros(size_bytes, dtype=torch.uint8, device=f"cuda:{gpu_id}")
ptr = tensor.data_ptr()
print(f" Allocated {size_bytes / 1e9:.1f} GB on GPU {gpu_id}")
return ptr, tensor # keep tensor alive
def run_target(args):
"""Run as target node: allocate KV cache pool and wait."""
from mooncake.engine import TransferEngine
print(f"=== Target Node ===")
print(f"Pool size: {args.pool_size_gb} GB")
print(f"Protocol: {args.protocol}")
engine = TransferEngine()
ret = engine.initialize(
args.local_server_name, args.metadata_server, args.protocol, ""
)
if ret != 0:
raise RuntimeError(f"Engine initialization failed: {ret}")
if args.metadata_server == "P2PHANDSHAKE":
host = args.local_server_name.rpartition(":")[0]
rpc_port = engine.get_rpc_port()
actual_name = f"{host}:{rpc_port}"
print(f"Actual server name: {actual_name}")
pool_bytes = int(args.pool_size_gb * 1024 * 1024 * 1024)
print(f"Allocating {args.pool_size_gb} GB KV cache pool...")
if args.use_gpu:
pool_addr, _tensor = allocate_gpu_memory(pool_bytes, args.gpu_id)
else:
pool_addr, _, _ = allocate_cpu_memory(pool_bytes)
print(f"Registering memory with transfer engine...")
t0 = time.time()
ret = engine.register_memory(pool_addr, pool_bytes)
reg_time = time.time() - t0
if ret != 0:
raise RuntimeError(f"Memory registration failed: {ret}")
print(f" Registration took {reg_time:.1f}s")
print(f"\nTarget ready. Pool addr: 0x{pool_addr:x}, size: {pool_bytes}")
print("Waiting for initiator (Ctrl+C to stop)...")
# Write pool info for initiator to read
info = {"pool_addr": pool_addr, "pool_bytes": pool_bytes}
print(f"TARGET_INFO:{json.dumps(info)}")
try:
signal.pause()
except KeyboardInterrupt:
print("\nShutting down target.")
def run_initiator(args):
"""Run as initiator node: pull prefix KV cache and measure performance."""
from mooncake.engine import TransferEngine
print(f"=== Initiator Node ===")
print(f"Target: {args.target_server_name}")
print(f"Pool size: {args.pool_size_gb} GB")
print(f"Protocol: {args.protocol}")
print(f"KV bytes/token: {args.kv_bytes_per_token}")
print(f"Threads: {args.threads}")
if not args.target_server_name:
raise RuntimeError("--target_server_name required in initiator mode")
prefix_tokens_list = [int(x) for x in args.prefix_tokens.split(",")]
transfer_sizes = [
tokens * args.kv_bytes_per_token for tokens in prefix_tokens_list
]
print(f"\nTest matrix:")
for tokens, size in zip(prefix_tokens_list, transfer_sizes):
print(f" {tokens:>6} tokens -> {size / 1e6:.1f} MB transfer")
# Initialize engine
engine = TransferEngine()
ret = engine.initialize(
args.local_server_name, args.metadata_server, args.protocol, ""
)
if ret != 0:
raise RuntimeError(f"Engine initialization failed: {ret}")
if args.metadata_server == "P2PHANDSHAKE":
host = args.local_server_name.rpartition(":")[0]
rpc_port = engine.get_rpc_port()
actual_name = f"{host}:{rpc_port}"
print(f"Actual server name: {actual_name}")
# Allocate local receive buffer (large enough for the biggest transfer)
max_transfer = max(transfer_sizes)
recv_bytes = max_transfer
print(f"\nAllocating {recv_bytes / 1e6:.1f} MB local receive buffer...")
if args.use_gpu:
recv_addr, _tensor = allocate_gpu_memory(recv_bytes, args.gpu_id)
else:
recv_addr, _, _ = allocate_cpu_memory(recv_bytes)
ret = engine.register_memory(recv_addr, recv_bytes)
if ret != 0:
raise RuntimeError(f"Local memory registration failed: {ret}")
# Get target's buffer address
print(f"Connecting to target {args.target_server_name}...")
remote_addr = engine.get_first_buffer_address(args.target_server_name)
if remote_addr == 0:
raise RuntimeError(
"Cannot get target buffer address. "
"Is the target running and registered?"
)
print(f" Remote buffer at 0x{remote_addr:x}")
# Connection warmup: a small transfer to establish the EFA connection
# (openSegment, endpoint creation, etc.) so it doesn't skew the first
# prefix size's measurements.
print("Warming up connection...")
warmup_size = min(transfer_sizes[0], recv_bytes)
for _ in range(3):
engine.transfer_sync_read(
args.target_server_name, recv_addr, remote_addr, warmup_size
)
print(" Connection ready.")
num_threads = args.threads
def do_transfer(local_addr, remote_addr_with_offset, size):
"""Single transfer_sync_read call, suitable for thread pool."""
return engine.transfer_sync_read(
args.target_server_name, local_addr,
remote_addr_with_offset, size
)
def threaded_transfer(local_base, remote_base, total_size, pool):
"""Split transfer across threads; return max of 0 (ok) or error code."""
if num_threads <= 1:
return engine.transfer_sync_read(
args.target_server_name, local_base, remote_base, total_size
)
chunk = total_size // num_threads
# Align chunk to 4KB
chunk = chunk & ~0xFFF
futures = []
for t in range(num_threads):
off = t * chunk
sz = chunk if t < num_threads - 1 else (total_size - off)
futures.append(
pool.submit(do_transfer, local_base + off,
remote_base + off, sz)
)
return max(f.result() for f in futures)
# Run benchmarks
print(f"\n{'='*72}")
print(
f"{'Prefix':>8} {'Size':>10} {'Latency(ms)':>12} "
f"{'p50(ms)':>10} {'p99(ms)':>10} {'Tput(GB/s)':>12}"
)
print(f"{'='*72}")
results = []
with ThreadPoolExecutor(max_workers=num_threads) as pool:
for tokens, transfer_size in zip(prefix_tokens_list, transfer_sizes):
if transfer_size > recv_bytes:
print(f" SKIP {tokens} tokens: transfer {transfer_size} > recv buffer")
continue
# Warmup: use the same offset pattern as the benchmark to pre-warm
# remote memory pages and DMA paths at each offset.
pool_bytes = int(args.pool_size_gb * 1024 * 1024 * 1024)
max_offset = pool_bytes - transfer_size
for w in range(args.warmup):
for i in range(args.iterations):
if max_offset > 0:
offset = ((i * transfer_size) % max_offset) & ~0xFFF
else:
offset = 0
ret = threaded_transfer(
recv_addr, remote_addr + offset, transfer_size, pool
)
if ret != 0:
print(f" WARNING: warmup transfer failed: {ret}")
break
# Benchmark
latencies = []
errors = 0
for i in range(args.iterations):
# Use different offset within remote pool for each iteration
# to simulate accessing different prefix locations
if max_offset > 0:
# Align to 4KB boundary
offset = ((i * transfer_size) % max_offset) & ~0xFFF
else:
offset = 0
t0 = time.perf_counter()
ret = threaded_transfer(
recv_addr, remote_addr + offset, transfer_size, pool
)
elapsed = time.perf_counter() - t0
if ret != 0:
errors += 1
if errors <= 3:
print(f" ERROR: transfer failed at iter {i}: {ret}")
continue
latencies.append(elapsed * 1000) # ms
if not latencies:
print(f" {tokens:>6}k ALL FAILED ({errors} errors)")
continue
latencies.sort()
avg_ms = statistics.mean(latencies)
p50_ms = latencies[len(latencies) // 2]
p99_ms = latencies[int(len(latencies) * 0.99)]
throughput_gbs = (transfer_size / 1e9) / (p50_ms / 1000)
print(
f" {tokens:>6} {transfer_size/1e6:>8.1f}MB "
f"{avg_ms:>11.2f} {p50_ms:>10.2f} {p99_ms:>10.2f} "
f"{throughput_gbs:>11.2f}"
)
results.append(
{
"prefix_tokens": tokens,
"transfer_bytes": transfer_size,
"transfer_mb": transfer_size / 1e6,
"pool_size_gb": args.pool_size_gb,
"iterations": len(latencies),
"errors": errors,
"avg_latency_ms": round(avg_ms, 3),
"p50_latency_ms": round(p50_ms, 3),
"p99_latency_ms": round(p99_ms, 3),
"throughput_gbs": round(throughput_gbs, 3),
"threads": num_threads,
}
)
print(f"{'='*72}")
# Summary
if results:
print(f"\n=== Summary (pool_size={args.pool_size_gb}GB, threads={num_threads}) ===")
print(json.dumps(results, indent=2))
# Save results
thread_tag = f"_t{num_threads}" if num_threads > 1 else ""
out_file = (
f"kvcache_bench_pool{args.pool_size_gb}gb"
f"{'_gpu' if args.use_gpu else '_cpu'}{thread_tag}.json"
)
with open(out_file, "w") as f:
json.dump(
{
"pool_size_gb": args.pool_size_gb,
"protocol": args.protocol,
"use_gpu": args.use_gpu,
"kv_bytes_per_token": args.kv_bytes_per_token,
"threads": num_threads,
"iterations": args.iterations,
"results": results,
},
f,
indent=2,
)
print(f"Results saved to {out_file}")
def main():
args = parse_args()
if args.mode == "target":
run_target(args)
else:
run_initiator(args)
if __name__ == "__main__":
main()

View File

@ -113,10 +113,6 @@ static inline int64_t getCurrentTimeInNano() {
return (int64_t{ts.tv_sec} * kNanosPerSecond + int64_t{ts.tv_nsec});
}
static inline int64_t getCurrentTimeInMilli() {
return getCurrentTimeInNano() / 1000 / 1000;
}
static inline std::string getCurrentDateTime() {
auto now = std::chrono::system_clock::now();
auto time_t_now = std::chrono::system_clock::to_time_t(now);
@ -331,13 +327,9 @@ static inline ssize_t writeFully(int fd, const void *buf, size_t len) {
}
static inline ssize_t readFully(int fd, void *buf, size_t len) {
// Set a timeout for read to avoid hanging forever.
constexpr std::chrono::seconds kReadTimeout = std::chrono::seconds(300);
const std::chrono::steady_clock::time_point deadline =
std::chrono::steady_clock::now() + kReadTimeout;
char *pos = (char *)buf;
size_t nbytes = len;
while (nbytes && std::chrono::steady_clock::now() < deadline) {
while (nbytes) {
ssize_t rc = read(fd, pos, nbytes);
if (rc < 0 && (errno == EAGAIN || errno == EINTR))
continue;
@ -352,14 +344,7 @@ static inline ssize_t readFully(int fd, void *buf, size_t len) {
pos += rc;
nbytes -= rc;
}
if (nbytes != 0) {
LOG(WARNING) << "Socket read timed out, timeout: "
<< kReadTimeout.count()
<< ", deadline: " << deadline.time_since_epoch().count()
<< ", read " << len - nbytes << " out of " << len
<< " bytes";
}
return len - nbytes;
return len;
}
static inline int writeString(int fd, const HandShakeRequestType type,
@ -472,21 +457,6 @@ static inline const std::string MakeNicPath(const std::string &server_name,
return server_name + NIC_PATH_DELIM + nic_name;
}
// Strip the port from a nic_path to get a stable key for endpoint reuse.
// "ip-172-31-45-191:15365@rdmap135s0" → "ip-172-31-45-191@rdmap135s0"
// This allows the same physical peer to reuse endpoints across reconnections
// (each run picks a random P2P handshake port).
static inline std::string normalizeNicPath(const std::string &nic_path) {
std::string server_name = getServerNameFromNicPath(nic_path);
std::string nic_name = getNicNameFromNicPath(nic_path);
if (server_name.empty() || nic_name.empty()) return nic_path;
size_t colon = server_name.rfind(':');
if (colon != std::string::npos) {
server_name = server_name.substr(0, colon);
}
return server_name + NIC_PATH_DELIM + nic_name;
}
static inline bool overlap(const void *a, size_t a_len, const void *b,
size_t b_len) {
return (a >= b && a < (char *)b + b_len) ||

View File

@ -65,7 +65,6 @@ struct GlobalConfig {
int ib_pci_relaxed_ordering_mode = 0;
bool ascend_use_fabric_mem = false;
bool ascend_agent_mode = false;
size_t efa_striping_threshold = 2 * 1024 * 1024; // 2MB default
// ub config parameters
size_t num_jfc_per_ctx = 2;
size_t num_jfce_per_ctx = 2;
@ -82,13 +81,13 @@ struct RpcCommunicatorConfig {
size_t pool_size = 10;
};
void loadGlobalConfig(GlobalConfig& config);
void loadGlobalConfig(GlobalConfig &config);
void dumpGlobalConfig();
void updateGlobalConfig(ibv_device_attr& device_attr);
void updateGlobalConfig(ibv_device_attr &device_attr);
GlobalConfig& globalConfig();
GlobalConfig &globalConfig();
uint16_t getDefaultHandshakePort();

View File

@ -55,14 +55,6 @@ class MultiTransport {
Transport *getTransport(const std::string &proto);
/**
* @brief Check if TCP is the only installed transport.
*
* When only TCP transport is available (no RDMA, NVLink, etc.),
* local memcpy is preferred over TCP loopback for same-host transfers.
*/
bool isTcpOnly() const;
std::vector<Transport *> listTransports();
void *getBaseAddr();

View File

@ -150,14 +150,6 @@ class TransferEngine {
Transport* getTransport(const std::string& proto);
/**
* @brief Check if TCP is the only installed transport.
*
* When only TCP transport is available (no RDMA, NVLink, etc.),
* local memcpy is preferred over TCP loopback for same-host transfers.
*/
bool isTcpOnly() const;
int syncSegmentCache(const std::string& segment_name = "");
std::shared_ptr<TransferMetadata> getMetadata();

View File

@ -105,8 +105,6 @@ transfer_engine_t createTransferEngine(const char *metadata_conn_string,
const char *ip_or_host_name,
uint64_t rpc_port, int auto_discover);
int discoverTopology(transfer_engine_t engine);
int getLocalIpAndPort(transfer_engine_t engine, char *buf_out, size_t buf_len);
transport_t installTransport(transfer_engine_t engine, const char *proto,
@ -121,11 +119,6 @@ segment_id_t openSegmentNoCache(transfer_engine_t engine,
int closeSegment(transfer_engine_t engine, segment_id_t segment_id);
// Eagerly pre-connect all EFA endpoints to `segment_name`. Eliminates the
// first-batch fi_av_insert stall (observed ~6 s for 16 local NICs × N peer
// NICs). No-op on non-EFA installs. Idempotent. Returns 0 on success.
int warmupEfaSegment(transfer_engine_t engine, const char *segment_name);
int removeLocalSegment(transfer_engine_t engine, const char *segment_name);
void destroyTransferEngine(transfer_engine_t engine);

View File

@ -334,8 +334,6 @@ class TransferEngineImpl {
return multi_transports_->getTransport(proto);
}
bool isTcpOnly() const { return multi_transports_->isTcpOnly(); }
int syncSegmentCache(const std::string& segment_name = "") {
return metadata_->syncSegmentCache(segment_name);
}

View File

@ -31,7 +31,6 @@
#include <memory>
#include <string>
#include <thread>
#include <map>
#include <unordered_map>
#include "common.h"
@ -45,87 +44,63 @@ class EfaTransport;
struct EfaCq {
EfaCq() : cq(nullptr), outstanding(0) {}
struct fid_cq* cq;
struct fid_cq *cq;
volatile int outstanding;
};
struct EfaMemoryRegionMeta {
void* addr;
void *addr;
size_t length;
struct fid_mr* mr;
struct fid_mr *mr;
uint64_t key;
};
// Endpoint store for EFA with LRU eviction support.
// When the store approaches the AV capacity (max_endpoints), stale endpoints
// that have been inactive beyond `inactive_timeout_sec` are evicted to free
// AV slots. This prevents AV exhaustion in long-running services that
// communicate with many transient peers.
// Simple endpoint store for EFA
class EfaEndpointStore {
public:
static constexpr double kDefaultInactiveTimeoutSec = 5.0; // 5 seconds
explicit EfaEndpointStore(
size_t max_endpoints = 65536,
double inactive_timeout_sec = kDefaultInactiveTimeoutSec);
std::shared_ptr<EfaEndPoint> get(const std::string& peer_nic_path);
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.
// Triggers eviction when the store is at or above max_endpoints.
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,
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);
void remove(const std::string &peer_nic_path);
int disconnectAll();
size_t size() const;
// Evict endpoints that have been inactive longer than the configured
// timeout. Returns the number of evicted endpoints.
size_t evictStale();
// Remove endpoints whose connections are broken (not connected and not
// initializing). Returns the number of removed endpoints.
size_t removeDisconnected();
private:
// Must be called with write lock held.
size_t evictStaleLocked();
mutable RWSpinlock lock_;
std::unordered_map<std::string, std::shared_ptr<EfaEndPoint>> endpoints_;
size_t max_endpoints_;
double inactive_timeout_sec_;
};
// 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(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 = 65536);
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
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);
int registerMemoryRegionInternal(void *addr, size_t length, int access,
EfaMemoryRegionMeta &mrMeta);
public:
bool active() const { return active_; }
@ -133,39 +108,27 @@ class EfaContext {
public:
// EndPoint Management
std::shared_ptr<EfaEndPoint> endpoint(const std::string& peer_nic_path);
int deleteEndpoint(const std::string& peer_nic_path);
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_; }
EfaTransport &engine() { return engine_; }
const EfaTransport &engine() const { return engine_; }
// Submit slices for transfer
int submitPostSend(const std::vector<Transport::Slice*>& slice_list);
int submitPostSend(const std::vector<Transport::Slice *> &slice_list);
// Poll completion queue for completed operations
int pollCq(int max_entries, int cq_index = 0);
// Evict stale endpoints to free EFA resources
size_t evictStaleEndpoints() {
return endpoint_store_ ? endpoint_store_->evictStale() : 0;
}
// Get CQ count
size_t cqCount() const { return cq_list_.size(); }
// Round-robin CQ assignment for new endpoints (mirrors RDMA transport)
std::shared_ptr<EfaCq> nextCq() {
if (cq_list_.empty()) return nullptr;
int index = (next_cq_index_++) % cq_list_.size();
return cq_list_[index];
}
// Get CQ outstanding count pointer
volatile int* cqOutstandingCount(int cq_index) {
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;
}
@ -179,10 +142,10 @@ class EfaContext {
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_; }
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)
@ -190,24 +153,23 @@ class EfaContext {
std::string gid() const { return localAddr(); }
private:
EfaTransport& engine_;
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
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_;
std::atomic<int> next_cq_index_{0};
RWSpinlock mr_lock_;
std::map<uint64_t, EfaMemoryRegionMeta> mr_map_;
std::unordered_map<uint64_t, EfaMemoryRegionMeta> mr_map_;
};
} // namespace mooncake

View File

@ -41,8 +41,8 @@ class EfaContext;
// 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
Transport::Slice *slice; // Slice pointer for completion handling
volatile int *wr_depth; // Pointer to endpoint's wr_depth_ for CQ
// completion decrement
};
@ -55,29 +55,28 @@ class EfaEndPoint {
enum Status { INITIALIZING, UNCONNECTED, CONNECTED };
EfaEndPoint(EfaContext& context);
EfaEndPoint(EfaContext &context);
~EfaEndPoint();
// Construct endpoint with specified completion queue
int construct(struct fid_cq* cq, volatile int* cq_outstanding,
size_t num_qp_list = 1, size_t max_sge = 4,
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);
void setPeerNicPath(const std::string &peer_nic_path);
int setupConnectionsByActive();
int setupConnectionsByActive(const std::string& peer_nic_path) {
int setupConnectionsByActive(const std::string &peer_nic_path) {
setPeerNicPath(peer_nic_path);
return setupConnectionsByActive();
}
int setupConnectionsByPassive(const HandShakeDesc& peer_desc,
HandShakeDesc& local_desc);
int setupConnectionsByPassive(const HandShakeDesc &peer_desc,
HandShakeDesc &local_desc);
bool hasOutstandingSlice() const;
@ -94,12 +93,6 @@ class EfaEndPoint {
return (getCurrentTimeInNano() - inactive_time_) / 1000000000.0;
}
void touchLastUsed() { last_used_time_ = getCurrentTimeInNano(); }
double lastUsedAge() const {
return (getCurrentTimeInNano() - last_used_time_) / 1000000000.0;
}
public:
bool connected() const {
return status_.load(std::memory_order_relaxed) == CONNECTED;
@ -115,8 +108,8 @@ class EfaEndPoint {
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);
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; }
@ -127,27 +120,27 @@ class EfaEndPoint {
// Get peer's fi_addr
fi_addr_t getPeerFiAddr() const { return peer_fi_addr_; }
EfaContext& context() { return context_; }
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);
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);
int insertPeerAddr(const std::string &peer_addr);
private:
EfaContext& context_;
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_;
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)
@ -156,15 +149,14 @@ class EfaEndPoint {
volatile int wr_depth_;
int max_wr_depth_;
volatile int* cq_outstanding_;
volatile int *cq_outstanding_;
// Spinlock to serialize fi_write/fi_read calls on this endpoint.
// 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_;
volatile uint64_t last_used_time_; // Updated on connection and I/O
};
} // namespace mooncake

View File

@ -52,96 +52,82 @@ class EfaTransport : public Transport {
~EfaTransport();
int install(std::string& local_server_name,
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"; }
const char *getName() const override { return "efa"; }
int registerLocalMemory(void* addr, size_t length,
const std::string& location, bool remote_accessible,
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 unregisterLocalMemory(void *addr, bool update_metadata = true) override;
int registerLocalMemoryBatch(const std::vector<BufferEntry>& buffer_list,
const std::string& location) override;
int registerLocalMemoryBatch(const std::vector<BufferEntry> &buffer_list,
const std::string &location) override;
int unregisterLocalMemoryBatch(
const std::vector<void*>& addr_list) override;
// Eagerly establish EFA endpoints to every NIC of `segment_name`.
//
// Rationale: libfabric FI_EP_RDM endpoints resolve peer addresses lazily
// via fi_av_insert() on first send. With 16 local NICs × N peer NICs,
// the first submitTransfer() of a batch serializes ~N*16
// handshake+fi_av_insert round-trips, producing a single-digit-second
// stall (observed ~6 s on B300 for the first 100 × 0.5 MB batch). After
// this call returns, every (local_ctx, peer_nic) endpoint is CONNECTED
// and the first real submitTransfer() goes straight to fi_write/fi_read.
//
// Safe to call multiple times (idempotent: endpoint() + setup are both
// idempotent). Re-run after any openSegment() on a new peer.
int warmupSegment(const std::string& segment_name);
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,
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,
int unregisterLocalMemoryInternal(void *addr, bool update_metadata,
bool force_sequential);
// TRANSFER
Status submitTransfer(BatchID batch_id,
const std::vector<TransferRequest>& entries) override;
const std::vector<TransferRequest> &entries) override;
Status submitTransferTask(
const std::vector<TransferTask*>& task_list) override;
const std::vector<TransferTask *> &task_list) override;
Status getTransferStatus(BatchID batch_id,
std::vector<TransferStatus>& status);
std::vector<TransferStatus> &status);
Status getTransferStatus(BatchID batch_id, size_t task_id,
TransferStatus& status) override;
TransferStatus &status) override;
SegmentID getSegmentID(const std::string& segment_name);
SegmentID getSegmentID(const std::string &segment_name);
private:
int allocateLocalSegmentID();
int preTouchMemory(void* addr, size_t length);
int preTouchMemory(void *addr, size_t length);
public:
int onSetupEfaConnections(const HandShakeDesc& peer_desc,
HandShakeDesc& local_desc);
int onSetupEfaConnections(const HandShakeDesc &peer_desc,
HandShakeDesc &local_desc);
int sendHandshake(const std::string& peer_server_name,
const HandShakeDesc& local_desc,
HandShakeDesc& peer_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_; }
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);
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);
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
@ -153,16 +139,6 @@ class EfaTransport : public Transport {
std::vector<std::shared_ptr<EfaContext>> context_list_;
std::shared_ptr<Topology> local_topology_;
// Track chunked MR registrations for per-NIC partitioned buffers.
// When a buffer exceeds max_mr_size, it is split into chunks, each
// registered on a disjoint subset of NICs (per-NIC partition).
struct ChunkRegistration {
uint64_t addr;
std::vector<size_t> nic_indices;
};
std::mutex chunk_map_mutex_;
std::unordered_map<uint64_t, std::vector<ChunkRegistration>> chunk_map_;
// CQ polling worker threads
std::atomic<bool> worker_running_{false};
std::vector<std::thread> worker_threads_;

View File

@ -31,7 +31,7 @@ class UbEndPoint {
using HandShakeDesc = TransferMetadata::HandShakeDesc;
UbEndPoint() : inactive_time_(0), active_(true), status_(INITIALIZING) {}
UbEndPoint() : status_(INITIALIZING), active_(true) {}
virtual int construct(GlobalConfig& config) = 0;

View File

@ -21,7 +21,7 @@
#include <utility>
#include "common.h"
#include "config.h"
#include "urma_api.h"
#include "ub/umdk/urma/urma_api.h"
#include "transport/kunpeng_transport/ub_context.h"
#include "transport/kunpeng_transport/ub_endpoint.h"

View File

@ -38,8 +38,6 @@ class EndpointStore {
public:
virtual std::shared_ptr<RdmaEndPoint> getEndpoint(
const std::string &peer_nic_path) = 0;
virtual std::shared_ptr<RdmaEndPoint> getEndpointByPtr(
const RdmaEndPoint *endpoint_ptr) = 0;
virtual std::shared_ptr<RdmaEndPoint> insertEndpoint(
const std::string &peer_nic_path, RdmaContext *context) = 0;
virtual int deleteEndpoint(const std::string &peer_nic_path) = 0;
@ -60,8 +58,6 @@ class FIFOEndpointStore : public EndpointStore {
FIFOEndpointStore(size_t max_size) : max_size_(max_size) {}
std::shared_ptr<RdmaEndPoint> getEndpoint(
const std::string &peer_nic_path) override;
std::shared_ptr<RdmaEndPoint> getEndpointByPtr(
const RdmaEndPoint *endpoint_ptr) override;
std::shared_ptr<RdmaEndPoint> insertEndpoint(
const std::string &peer_nic_path, RdmaContext *context) override;
int deleteEndpoint(const std::string &peer_nic_path) override;
@ -93,8 +89,6 @@ class SIEVEEndpointStore : public EndpointStore {
: waiting_list_len_(0), max_size_(max_size) {}
std::shared_ptr<RdmaEndPoint> getEndpoint(
const std::string &peer_nic_path) override;
std::shared_ptr<RdmaEndPoint> getEndpointByPtr(
const RdmaEndPoint *endpoint_ptr) override;
std::shared_ptr<RdmaEndPoint> insertEndpoint(
const std::string &peer_nic_path, RdmaContext *context) override;
int deleteEndpoint(const std::string &peer_nic_path) override;

View File

@ -99,9 +99,6 @@ class RdmaContext {
// EndPoint Management
std::shared_ptr<RdmaEndPoint> endpoint(const std::string &peer_nic_path);
std::shared_ptr<RdmaEndPoint> getEndpointByPtr(
const RdmaEndPoint *endpoint_ptr);
int deleteEndpoint(const std::string &peer_nic_path);
int disconnectAllEndpoints();

View File

@ -1 +0,0 @@
src/customer_pattern.rs

View File

@ -17,84 +17,16 @@ use std::path::PathBuf;
fn main() {
println!("cargo:rustc-link-search=native=../build/src");
println!("cargo:rustc-link-search=native=../../build/mooncake-transfer-engine/src");
println!("cargo:rustc-link-lib=static=transfer_engine");
// libbase.a holds mooncake::Status, which libtransfer_engine.a references.
println!("cargo:rustc-link-search=native=../build/src/common/base");
println!("cargo:rustc-link-search=native=../../build/mooncake-transfer-engine/src/common/base");
println!("cargo:rustc-link-lib=static=base");
// The transfer_engine build uses ASIO_SEPARATE_COMPILATION + ASIO_DYN_LINK,
// so the asio symbols live in mooncake-asio/libasio.so. Link it whenever
// we can find it (standalone cmake build places it alongside src/).
println!("cargo:rustc-link-search=native=../build/mooncake-asio");
println!("cargo:rustc-link-search=native=../../build/mooncake-asio");
println!("cargo:rustc-link-lib=asio");
// EFA on AWS installs libfabric under /opt/amazon/efa/lib.
if std::path::Path::new("/opt/amazon/efa/lib").exists() {
println!("cargo:rustc-link-search=native=/opt/amazon/efa/lib");
}
println!("cargo:rustc-link-lib=stdc++");
println!("cargo:rustc-link-lib=ibverbs");
// libfabric (fi_*): only needed for EFA transport, but harmless when
// the system has it installed; required on AWS EFA hosts. Opt-out by
// setting MOONCAKE_WITHOUT_LIBFABRIC=1 if building on a box without it.
if env::var("MOONCAKE_WITHOUT_LIBFABRIC")
.map(|v| v == "1" || v.eq_ignore_ascii_case("on") || v.eq_ignore_ascii_case("true"))
.unwrap_or(false)
{
// skip
} else {
println!("cargo:rustc-link-lib=fabric");
}
println!("cargo:rustc-link-lib=glog");
println!("cargo:rustc-link-lib=gflags");
println!("cargo:rustc-link-lib=pthread");
println!("cargo:rustc-link-lib=jsoncpp");
println!("cargo:rustc-link-lib=numa");
println!("cargo:rustc-link-lib=curl");
let flag_on = |name: &str| {
env::var(name)
.map(|v| v == "1" || v.eq_ignore_ascii_case("on") || v.eq_ignore_ascii_case("true"))
.unwrap_or(false)
};
// etcd-cpp-api: only needed when transfer_engine was built with
// USE_ETCD=ON. Opt-in via MOONCAKE_WITH_ETCD=1 to keep non-etcd builds
// (e.g. EFA-only on AWS) linkable.
if flag_on("MOONCAKE_WITH_ETCD") {
println!("cargo:rustc-link-lib=etcd-cpp-api");
}
// CUDA runtime: libtransfer_engine.a built with USE_CUDA=ON pulls in
// cudaMemcpy/cudaStream* symbols. The Rust demos themselves don't call
// CUDA — this is purely a transitive archive dep. Enable with
// MOONCAKE_WITH_CUDA=1 and optional CUDA_HOME override for lib path.
if flag_on("MOONCAKE_WITH_CUDA") {
// Accept either a CUDA_HOME (append lib64/lib) or an explicit
// CUDART_LIB_DIR that already points at the directory containing
// libcudart.so. This covers both /usr/local/cuda installs and
// pip-wheel layouts like .../nvidia/cu13/lib.
if let Ok(dir) = env::var("CUDART_LIB_DIR") {
println!("cargo:rustc-link-search=native={}", dir);
} else if let Ok(cuda_home) = env::var("CUDA_HOME") {
let lib64 = PathBuf::from(&cuda_home).join("lib64");
let lib = PathBuf::from(&cuda_home).join("lib");
if lib64.exists() {
println!("cargo:rustc-link-search=native={}", lib64.display());
}
if lib.exists() {
println!("cargo:rustc-link-search=native={}", lib.display());
}
} else {
println!("cargo:rustc-link-search=native=/usr/local/cuda/lib64");
}
println!("cargo:rustc-link-lib=cudart");
}
println!("cargo:rustc-link-lib=etcd-cpp-api");
let bindings = bindgen::builder()
.header("../include/transfer_engine_c.h")

View File

@ -91,27 +91,6 @@ impl TransferEngine {
Ok(Self { engine })
}
pub fn discover_topology(&self) -> Result<()> {
let ret = unsafe { bindings::discoverTopology(self.engine) };
if ret != 0 {
bail!("Failed to discover topology")
} else {
Ok(())
}
}
pub fn install_transport(&self, proto: &str) -> Result<()> {
let proto_c = CString::new(proto).map_err(|_| anyhow!("CString::new failed"))?;
let ret = unsafe {
bindings::installTransport(self.engine, proto_c.as_ptr(), std::ptr::null_mut())
};
if ret.is_null() {
bail!("Failed to install transport '{}'", proto)
} else {
Ok(())
}
}
pub fn close(&mut self) -> Result<()> {
unsafe {
bindings::destroyTransferEngine(self.engine);
@ -263,20 +242,6 @@ impl TransferEngine {
}
}
/// Eagerly establish EFA endpoints to `segment_name` so the first
/// `submit_transfer` doesn't pay the serial fi_av_insert cost. No-op on
/// non-EFA transports. Call after `open_segment` (and after the metadata
/// has the peer's NIC list published).
pub fn warmup_efa_segment(&self, name: &str) -> Result<()> {
let name_c = CString::new(name).map_err(|_| anyhow!("CString::new failed"))?;
let ret = unsafe { bindings::warmupEfaSegment(self.engine, name_c.as_ptr()) };
if ret < 0 {
bail!("warmupEfaSegment failed for {}: {}", name, ret)
} else {
Ok(())
}
}
pub fn sync_segment_cache(&self) -> Result<()> {
let ret = unsafe { bindings::syncSegmentCache(self.engine) };
if ret < 0 {

View File

@ -70,10 +70,15 @@ if(USE_CUDA)
endif()
if(USE_MACA)
if(DEFINED ENV{MACA_HOME})
set(MACA_HOME $ENV{MACA_HOME})
else()
set(MACA_HOME /opt/maca)
endif()
if(NOT DEFINED MACA_RUNTIME_LIBS)
set(MACA_RUNTIME_LIBS mcruntime mxc-runtime64 rt)
endif()
target_include_directories(transfer_engine PRIVATE ${MACA_INCLUDE_DIR})
target_include_directories(transfer_engine PRIVATE ${MACA_HOME}/include)
target_link_libraries(transfer_engine PUBLIC ${MACA_RUNTIME_LIBS})
endif()
@ -91,22 +96,20 @@ if(USE_MLU)
target_link_libraries(transfer_engine PUBLIC cnrt cndrv)
endif()
if(USE_ASCEND)
target_link_libraries(transfer_engine PUBLIC ascendcl hccl ascend_transport
MPI::MPI)
if (USE_ASCEND)
target_link_libraries(transfer_engine PUBLIC ascendcl hccl ascend_transport MPI::MPI)
endif()
if(USE_ASCEND_DIRECT)
if (USE_ASCEND_DIRECT)
target_link_libraries(transfer_engine PUBLIC ascend_transport)
endif()
if(USE_UBSHMEM)
if (USE_UBSHMEM)
target_link_libraries(transfer_engine PUBLIC ascend_transport)
endif()
if(USE_ASCEND_HETEROGENEOUS)
file(GLOB ASCEND_TOOLKIT_ROOT
"/usr/local/Ascend/ascend-toolkit/latest/*-linux")
if (USE_ASCEND_HETEROGENEOUS)
file(GLOB ASCEND_TOOLKIT_ROOT "/usr/local/Ascend/ascend-toolkit/latest/*-linux")
set(ASCEND_LIB_DIR "${ASCEND_TOOLKIT_ROOT}/lib64")
link_directories(${ASCEND_LIB_DIR})
target_link_libraries(transfer_engine PUBLIC ascendcl ascend_transport)
@ -114,7 +117,7 @@ endif()
if(USE_TENT)
add_compile_definitions(transfer_engine PUBLIC USE_TENT)
target_link_libraries(transfer_engine PUBLIC tent_link_group)
target_link_libraries(transfer_engine PUBLIC tent)
endif()
if(USE_INTRA_NVLINK)

View File

@ -21,8 +21,8 @@
#include <unistd.h>
namespace mooncake {
void loadGlobalConfig(GlobalConfig& config) {
const char* num_cq_per_ctx_env = std::getenv("MC_NUM_CQ_PER_CTX");
void loadGlobalConfig(GlobalConfig &config) {
const char *num_cq_per_ctx_env = std::getenv("MC_NUM_CQ_PER_CTX");
if (num_cq_per_ctx_env) {
int val = atoi(num_cq_per_ctx_env);
if (val > 0 && val < 256) {
@ -35,7 +35,7 @@ void loadGlobalConfig(GlobalConfig& config) {
<< "Ignore value from environment variable MC_NUM_CQ_PER_CTX";
}
const char* num_comp_channels_per_ctx_env =
const char *num_comp_channels_per_ctx_env =
std::getenv("MC_NUM_COMP_CHANNELS_PER_CTX");
if (num_comp_channels_per_ctx_env) {
int val = atoi(num_comp_channels_per_ctx_env);
@ -46,7 +46,7 @@ void loadGlobalConfig(GlobalConfig& config) {
"MC_NUM_COMP_CHANNELS_PER_CTX";
}
const char* port_env = std::getenv("MC_IB_PORT");
const char *port_env = std::getenv("MC_IB_PORT");
if (port_env) {
int val = atoi(port_env);
if (val >= 0 && val < 256)
@ -55,7 +55,7 @@ void loadGlobalConfig(GlobalConfig& config) {
LOG(WARNING) << "Ignore value from environment variable MC_IB_PORT";
}
const char* gid_index_env = std::getenv("MC_GID_INDEX");
const char *gid_index_env = std::getenv("MC_GID_INDEX");
if (!gid_index_env) gid_index_env = std::getenv("NCCL_IB_GID_INDEX");
if (gid_index_env) {
@ -67,7 +67,7 @@ void loadGlobalConfig(GlobalConfig& config) {
<< "Ignore value from environment variable MC_GID_INDEX";
}
const char* max_cqe_per_ctx_env = std::getenv("MC_MAX_CQE_PER_CTX");
const char *max_cqe_per_ctx_env = std::getenv("MC_MAX_CQE_PER_CTX");
if (max_cqe_per_ctx_env) {
size_t val = atoi(max_cqe_per_ctx_env);
if (val > 0 && val <= UINT16_MAX)
@ -77,7 +77,7 @@ void loadGlobalConfig(GlobalConfig& config) {
<< "Ignore value from environment variable MC_MAX_CQE_PER_CTX";
}
const char* max_ep_per_ctx_env = std::getenv("MC_MAX_EP_PER_CTX");
const char *max_ep_per_ctx_env = std::getenv("MC_MAX_EP_PER_CTX");
if (max_ep_per_ctx_env) {
size_t val = atoi(max_ep_per_ctx_env);
if (val > 0 && val <= UINT16_MAX)
@ -87,7 +87,7 @@ void loadGlobalConfig(GlobalConfig& config) {
<< "Ignore value from environment variable MC_MAX_EP_PER_CTX";
}
const char* num_qp_per_ep_env = std::getenv("MC_NUM_QP_PER_EP");
const char *num_qp_per_ep_env = std::getenv("MC_NUM_QP_PER_EP");
if (num_qp_per_ep_env) {
int val = atoi(num_qp_per_ep_env);
if (val > 0 && val < 256)
@ -97,7 +97,7 @@ void loadGlobalConfig(GlobalConfig& config) {
<< "Ignore value from environment variable MC_NUM_QP_PER_EP";
}
const char* max_sge_env = std::getenv("MC_MAX_SGE");
const char *max_sge_env = std::getenv("MC_MAX_SGE");
if (max_sge_env) {
size_t val = atoi(max_sge_env);
if (val > 0 && val <= UINT16_MAX)
@ -106,7 +106,7 @@ void loadGlobalConfig(GlobalConfig& config) {
LOG(WARNING) << "Ignore value from environment variable MC_MAX_SGE";
}
const char* max_wr_env = std::getenv("MC_MAX_WR");
const char *max_wr_env = std::getenv("MC_MAX_WR");
if (max_wr_env) {
size_t val = atoi(max_wr_env);
if (val > 0 && val <= UINT16_MAX)
@ -115,7 +115,7 @@ void loadGlobalConfig(GlobalConfig& config) {
LOG(WARNING) << "Ignore value from environment variable MC_MAX_WR";
}
const char* max_inline_env = std::getenv("MC_MAX_INLINE");
const char *max_inline_env = std::getenv("MC_MAX_INLINE");
if (max_inline_env) {
size_t val = atoi(max_inline_env);
if (val <= UINT16_MAX)
@ -125,7 +125,7 @@ void loadGlobalConfig(GlobalConfig& config) {
<< "Ignore value from environment variable MC_MAX_INLINE";
}
const char* mtu_length_env = std::getenv("MC_MTU");
const char *mtu_length_env = std::getenv("MC_MTU");
if (mtu_length_env) {
size_t val = atoi(mtu_length_env);
if (val == 512)
@ -143,7 +143,7 @@ void loadGlobalConfig(GlobalConfig& config) {
}
}
const char* handshake_port_env = std::getenv("MC_HANDSHAKE_PORT");
const char *handshake_port_env = std::getenv("MC_HANDSHAKE_PORT");
if (handshake_port_env) {
int val = atoi(handshake_port_env);
if (val > 0 && val < 65536)
@ -153,7 +153,7 @@ void loadGlobalConfig(GlobalConfig& config) {
<< "Ignore value from environment variable MC_HANDSHAKE_PORT";
}
const char* workers_per_ctx_env = std::getenv("MC_WORKERS_PER_CTX");
const char *workers_per_ctx_env = std::getenv("MC_WORKERS_PER_CTX");
if (workers_per_ctx_env) {
size_t val = atoi(workers_per_ctx_env);
if (val > 0 && val <= 8)
@ -163,7 +163,7 @@ void loadGlobalConfig(GlobalConfig& config) {
<< "Ignore value from environment variable MC_WORKERS_PER_CTX";
}
const char* slice_size_env = std::getenv("MC_SLICE_SIZE");
const char *slice_size_env = std::getenv("MC_SLICE_SIZE");
if (slice_size_env) {
size_t val = atoi(slice_size_env);
if (val > 0)
@ -173,7 +173,7 @@ void loadGlobalConfig(GlobalConfig& config) {
<< "Ignore value from environment variable MC_SLICE_SIZE";
}
const char* min_reg_size_env = std::getenv("MC_MIN_REG_SIZE");
const char *min_reg_size_env = std::getenv("MC_MIN_REG_SIZE");
if (min_reg_size_env) {
size_t val = atoll(min_reg_size_env);
if (val > 0) {
@ -184,7 +184,7 @@ void loadGlobalConfig(GlobalConfig& config) {
<< "Ignore value from environment variable MC_MIN_REG_SIZE";
}
const char* max_mr_size_env = std::getenv("MC_MAX_MR_SIZE");
const char *max_mr_size_env = std::getenv("MC_MAX_MR_SIZE");
if (max_mr_size_env) {
uint64_t val = atoll(max_mr_size_env);
if (val > 0) {
@ -195,17 +195,7 @@ void loadGlobalConfig(GlobalConfig& config) {
}
}
const char* efa_striping_env = std::getenv("MC_EFA_STRIPING_THRESHOLD");
if (efa_striping_env) {
size_t val = strtoull(efa_striping_env, nullptr, 10);
if (val > 0)
config.efa_striping_threshold = val;
else
LOG(WARNING) << "Ignore value from environment variable "
"MC_EFA_STRIPING_THRESHOLD";
}
const char* retry_cnt_env = std::getenv("MC_RETRY_CNT");
const char *retry_cnt_env = std::getenv("MC_RETRY_CNT");
if (retry_cnt_env) {
size_t val = atoi(retry_cnt_env);
if (val > 0 && val < 128)
@ -215,12 +205,12 @@ void loadGlobalConfig(GlobalConfig& config) {
<< "Ignore value from environment variable MC_RETRY_CNT";
}
const char* disable_metacache = std::getenv("MC_DISABLE_METACACHE");
const char *disable_metacache = std::getenv("MC_DISABLE_METACACHE");
if (disable_metacache) {
config.metacache = false;
}
const char* handshake_listen_backlog =
const char *handshake_listen_backlog =
std::getenv("MC_HANDSHAKE_LISTEN_BACKLOG");
if (handshake_listen_backlog) {
int val = std::stoi(handshake_listen_backlog);
@ -232,7 +222,7 @@ void loadGlobalConfig(GlobalConfig& config) {
}
}
const char* log_level = std::getenv("MC_LOG_LEVEL");
const char *log_level = std::getenv("MC_LOG_LEVEL");
config.trace = false;
if (log_level) {
if (strcmp(log_level, "TRACE") == 0) {
@ -248,7 +238,7 @@ void loadGlobalConfig(GlobalConfig& config) {
}
FLAGS_minloglevel = config.log_level;
const char* slice_timeout_env = std::getenv("MC_SLICE_TIMEOUT");
const char *slice_timeout_env = std::getenv("MC_SLICE_TIMEOUT");
if (slice_timeout_env) {
int val = atoi(slice_timeout_env);
if (val > 0 && val < 65536)
@ -258,7 +248,7 @@ void loadGlobalConfig(GlobalConfig& config) {
<< "Ignore value from environment variable MC_SLICE_TIMEOUT";
}
const char* log_dir_path = std::getenv("MC_LOG_DIR");
const char *log_dir_path = std::getenv("MC_LOG_DIR");
if (log_dir_path) {
google::InitGoogleLogging("mooncake-transfer-engine");
if (opendir(log_dir_path) == NULL) {
@ -277,7 +267,7 @@ void loadGlobalConfig(GlobalConfig& config) {
}
}
const char* min_port_env = std::getenv("MC_MIN_PRC_PORT");
const char *min_port_env = std::getenv("MC_MIN_PRC_PORT");
if (min_port_env) {
int val = atoi(min_port_env);
if (val > 0 && val < 65536)
@ -287,7 +277,7 @@ void loadGlobalConfig(GlobalConfig& config) {
<< "Ignore value from environment variable MC_PRC_MIN_PORT";
}
const char* max_port_env = std::getenv("MC_MAX_PRC_PORT");
const char *max_port_env = std::getenv("MC_MAX_PRC_PORT");
if (max_port_env) {
int val = atoi(max_port_env);
if (val > 0 && val < 65536)
@ -301,7 +291,7 @@ void loadGlobalConfig(GlobalConfig& config) {
config.use_ipv6 = true;
}
const char* fragment_ratio = std::getenv("MC_FRAGMENT_RATIO");
const char *fragment_ratio = std::getenv("MC_FRAGMENT_RATIO");
if (fragment_ratio) {
size_t val = atoi(fragment_ratio);
if (val > 0 && val < config.slice_size)
@ -317,7 +307,7 @@ void loadGlobalConfig(GlobalConfig& config) {
config.enable_dest_device_affinity = true;
}
const char* enable_parallel_reg_mr =
const char *enable_parallel_reg_mr =
std::getenv("MC_ENABLE_PARALLEL_REG_MR");
if (enable_parallel_reg_mr) {
int val = atoi(enable_parallel_reg_mr);
@ -329,7 +319,7 @@ void loadGlobalConfig(GlobalConfig& config) {
}
}
const char* endpoint_store_type_env = std::getenv("MC_ENDPOINT_STORE_TYPE");
const char *endpoint_store_type_env = std::getenv("MC_ENDPOINT_STORE_TYPE");
if (endpoint_store_type_env) {
if (strcmp(endpoint_store_type_env, "FIFO") == 0) {
config.endpoint_store_type = EndpointStoreType::FIFO;
@ -341,7 +331,7 @@ void loadGlobalConfig(GlobalConfig& config) {
}
}
const char* traffic_class_env = std::getenv("MC_IB_TC");
const char *traffic_class_env = std::getenv("MC_IB_TC");
if (traffic_class_env) {
try {
int val = std::stoi(traffic_class_env);
@ -353,13 +343,13 @@ void loadGlobalConfig(GlobalConfig& config) {
<< "value " << traffic_class_env
<< " out of range (should be 0-255)";
}
} catch (const std::exception& e) {
} catch (const std::exception &e) {
LOG(WARNING) << "Invalid MC_IB_TC environment value: "
<< traffic_class_env << ". Error: " << e.what();
}
}
const char* ib_relaxed_ordering_env =
const char *ib_relaxed_ordering_env =
std::getenv("MC_IB_PCI_RELAXED_ORDERING");
if (ib_relaxed_ordering_env) {
int val = atoi(ib_relaxed_ordering_env);
@ -384,8 +374,8 @@ std::string mtuLengthToString(ibv_mtu mtu) {
return "UNKNOWN";
}
void updateGlobalConfig(ibv_device_attr& device_attr) {
auto& config = globalConfig();
void updateGlobalConfig(ibv_device_attr &device_attr) {
auto &config = globalConfig();
if (config.max_ep_per_ctx * config.num_qp_per_ep >
(size_t)device_attr.max_qp)
config.max_ep_per_ctx = device_attr.max_qp / config.num_qp_per_ep;
@ -402,7 +392,7 @@ void updateGlobalConfig(ibv_device_attr& device_attr) {
}
void dumpGlobalConfig() {
auto& config = globalConfig();
auto &config = globalConfig();
LOG(INFO) << "=== GlobalConfig ===";
LOG(INFO) << "num_cq_per_ctx = " << config.num_cq_per_ctx;
LOG(INFO) << "num_comp_channels_per_ctx = "
@ -421,7 +411,7 @@ void dumpGlobalConfig() {
LOG(INFO) << "ib_traffic_class = " << config.ib_traffic_class;
}
GlobalConfig& globalConfig() {
GlobalConfig &globalConfig() {
static GlobalConfig config;
static std::once_flag g_once_flag;
std::call_once(g_once_flag, []() { loadGlobalConfig(config); });

View File

@ -461,10 +461,6 @@ Transport* MultiTransport::getTransport(const std::string& proto) {
return transport_map_[proto].get();
}
bool MultiTransport::isTcpOnly() const {
return transport_map_.size() == 1 && transport_map_.count("tcp") == 1;
}
std::vector<Transport*> MultiTransport::listTransports() {
std::vector<Transport*> transport_list;
for (auto& entry : transport_map_)

View File

@ -27,6 +27,7 @@
#include <limits.h>
#include <stdlib.h>
#include <string.h>
#include <libgen.h>
#include <sys/types.h>
#include <sys/stat.h>
@ -34,8 +35,7 @@
#include "memory_location.h"
#include "topology.h"
#ifdef USE_UB
#include <libgen.h>
#include <urma_api.h>
#include <ub/umdk/urma/urma_api.h>
#endif
namespace mooncake {
@ -208,7 +208,9 @@ static std::vector<UBDevice> listUBDevices(
int num_devices = 0;
std::vector<UBDevice> devices;
urma_init_attr_t init_attr = {};
urma_init_attr_t init_attr = {
.uasid = 0,
};
if (urma_init(&init_attr) != URMA_SUCCESS) {
LOG(WARNING) << "Failed to urma init";
return {};
@ -472,21 +474,22 @@ void Topology::clear() {
int Topology::discover(const std::vector<std::string> &filter) {
matrix_.clear();
#ifdef USE_UB
auto all_hca = listUBDevices(filter);
for (auto &ent : discoverCpuTopology(all_hca)) {
matrix_[ent.name] = ent;
}
#else
auto all_hca = listInfiniBandDevices(filter);
for (auto &ent : discoverCpuTopology(all_hca)) {
matrix_[ent.name] = ent;
}
#endif
#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
defined(USE_MLU) || defined(USE_MACA)
for (auto &ent : discoverCudaTopology(all_hca)) {
matrix_[ent.name] = ent;
}
#endif
#ifdef USE_UB
auto ub_all_hca = listUBDevices(filter);
for (auto &ent : discoverCpuTopology(ub_all_hca)) {
matrix_[ent.name] = ent;
}
#endif
return resolve();
}

View File

@ -179,8 +179,6 @@ Transport* TransferEngine::getTransport(const std::string& proto) {
return impl_->getTransport(proto);
}
bool TransferEngine::isTcpOnly() const { return impl_->isTcpOnly(); }
int TransferEngine::syncSegmentCache(const std::string& segment_name) {
return impl_->syncSegmentCache(segment_name);
}
@ -580,15 +578,6 @@ Transport* TransferEngine::getTransport(const std::string& proto) {
return impl_->getTransport(proto);
}
bool TransferEngine::isTcpOnly() const {
if (use_tent_)
// TENT already rejects TCP loopback transfers when MC_STORE_MEMCPY
// is disabled, so auto-enabling memcpy is unnecessary in TENT mode.
return false;
else
return impl_->isTcpOnly();
}
int TransferEngine::syncSegmentCache(const std::string& segment_name) {
if (use_tent_)
return 0;

View File

@ -19,9 +19,6 @@
#include "transfer_engine.h"
#include "transport/transport.h"
#ifdef USE_EFA
#include "transport/efa_transport/efa_transport.h"
#endif
using namespace mooncake;
@ -39,11 +36,6 @@ transfer_engine_t createTransferEngine(const char *metadata_conn_string,
return (transfer_engine_t)native;
}
int discoverTopology(transfer_engine_t engine) {
TransferEngine *native = (TransferEngine *)engine;
return native->getLocalTopology()->discover({});
}
int getLocalIpAndPort(transfer_engine_t engine, char *buf_out, size_t buf_len) {
TransferEngine *native = (TransferEngine *)engine;
auto str = native->getLocalIpAndPort();
@ -85,21 +77,6 @@ int closeSegment(transfer_engine_t engine, segment_id_t segment_id) {
return native->closeSegment(segment_id);
}
int warmupEfaSegment(transfer_engine_t engine, const char *segment_name) {
#ifdef USE_EFA
TransferEngine *native = (TransferEngine *)engine;
auto *t = native->getTransport("efa");
if (!t) return 0; // Non-EFA build or EFA not installed; nothing to do.
auto *efa = dynamic_cast<EfaTransport *>(t);
if (!efa) return 0;
return efa->warmupSegment(segment_name ? segment_name : "");
#else
(void)engine;
(void)segment_name;
return 0;
#endif
}
int removeLocalSegment(transfer_engine_t engine, const char *segment_name) {
TransferEngine *native = (TransferEngine *)engine;
return native->removeLocalSegment(segment_name);

View File

@ -29,9 +29,7 @@
#include "transfer_metadata_plugin.h"
#include "transport/transport.h"
#ifdef USE_BAREX
#include "transport/barex_transport/barex_transport.h"
#endif
namespace mooncake {

Some files were not shown because too many files have changed in this diff Show More