Compare commits

..

12 Commits

Author SHA1 Message Date
yejj710 a6bfb2e195 do not use lru list defautly & add log for remove-event 2026-04-22 11:34:22 +08:00
yejj710 c837624931 Add LRU eviction policy to the index data of Mooncake-Store 2026-04-19 14:49:08 +08:00
yejj c67649c83a
support remove event & add some info-level log (#1881) 2026-04-13 21:09:18 +08:00
yejj 6ab4e40ee5
bugfix: Support multiple identical instances under a single tenant (#1840) 2026-04-08 15:43:51 +08:00
Asher Zhang 80c8b41cad
update Mooncake Store KVEvent Publisher (#1638) 2026-03-10 19:25:44 +08:00
Asher Zhang 770031b101
init publisher (#1635) 2026-03-09 20:25:49 +08:00
yejj710 fb7313b87d Merge remote-tracking branch 'center/main' into sync_code_from_master 2026-03-09 10:51:21 +08:00
yejj c15a911ae0
support mooncake_conductor global kv-indexer (#1614) 2026-03-05 20:26:51 +08:00
Ryan·Chang 45402ff9a1
[Misc] feat: support dynamic register and unregister API for KV Event Manager (#1569)
* Register/Unregister

* Register/Unregister

* add Register/Unregister interface

* fix: address review comments

* feat: replace LoraID with LoraName and optimize unregister
2026-02-28 21:21:57 +08:00
Ryan·Chang fbff581166
[Misc] Configure slog log level via environment variable (#1529)
* log level

* Delete .trae directory

Delete additional files.

---------

Co-authored-by: Ryanic-Chang <ryanchang1412@gmail.com>
2026-02-10 16:59:58 +08:00
yejj 0cc4fea731
Merge pull request #1478 from yejj710/kv-indexer-dev
[Misc][Dev] add indexer demo version & using example
2026-02-06 17:40:32 +08:00
Liziqi-77 d0ec8173d3 add indexer demo & using example 2026-02-03 10:23:48 +08:00
517 changed files with 17952 additions and 70501 deletions

View File

@ -1,266 +0,0 @@
---
name: mooncake-ci-local
description: Run Mooncake CI test suite locally — maps GitHub Actions CI steps to local commands. Use this skill whenever the user wants to run tests locally, reproduce a CI failure, check if their changes break tests, or run any subset of the CI test suite (C++ unit tests via ctest, Python integration tests, code format checks, or the full test pipeline). Trigger on phrases like "run tests", "run CI locally", "reproduce CI failure", "check my changes", "test before PR", "run ctest", "run python tests", "run all tests".
---
# Mooncake CI Local Test Runner
You help users run the Mooncake CI test suite locally. The CI has three test layers. Map what the user wants to the right layer, check prerequisites, and run the tests.
## CI Test Layers
### Layer 1 — C++ Unit Tests (ctest)
**CI equivalent:** `build` job in `ci.yml` — "Test (in build env) with coverage"
**Prerequisite services:**
```bash
# 1. etcd (port 2379)
etcd --advertise-client-urls http://127.0.0.1:2379 --listen-client-urls http://127.0.0.1:2379 &
sleep 2
etcdctl --endpoints=http://127.0.0.1:2379 endpoint health # verify
# 2. HTTP metadata server (port 8080)
cd mooncake-transfer-engine/example/http-metadata-server-python
pip install aiohttp
python ./bootstrap_server.py &
cd -
```
**Run:**
```bash
cd build
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 ctest -j --output-on-failure
```
**Run specific test:**
```bash
cd build
MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 ctest -R <test_name_pattern> --output-on-failure
# List all available tests: ctest -N
```
### Layer 2 — Python Integration Tests
**CI equivalent:** `test-wheel-ubuntu` job — `run_tests.sh`
**Prerequisite:** Mooncake wheel must be installed (either via `pip install` or via `make install` after build).
**Check install:**
```bash
python -c "import mooncake; print('OK')"
which mooncake_master # must NOT be /usr/local/bin (must be from Python package)
```
**Run full suite:**
```bash
# Start metadata server first
mooncake_http_metadata_server --port 8080 &
sleep 1
cd mooncake-wheel/tests
MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 MC_FORCE_TCP=true \
bash ../../scripts/run_tests.sh
```
**Individual Python tests** (all require metadata server + mooncake_master on port 50051):
```bash
# Setup shared services
mooncake_http_metadata_server --port 8080 &
mooncake_master --default_kv_lease_ttl=500 &
sleep 2
cd mooncake-wheel/tests
export MC_METADATA_SERVER=http://127.0.0.1:8080/metadata
export DEFAULT_KV_LEASE_TTL=500
export MC_FORCE_TCP=true
# Pick any test:
python test_distributed_object_store.py
python test_replicated_distributed_object_store.py
python test_put_get_tensor.py # requires torch + numpy
python test_safetensor_functions.py # requires safetensors
python test_dummy_client.py
python test_cli.py
python test_distributed_object_store_cxl.py # requires CXL build
```
**Transfer engine tests specifically:**
```bash
cd mooncake-wheel/tests
MC_METADATA_SERVER=http://127.0.0.1:8080/metadata MC_FORCE_TCP=true python transfer_engine_target.py &
TARGET_PID=$!
MC_METADATA_SERVER=http://127.0.0.1:8080/metadata MC_FORCE_TCP=true python transfer_engine_initiator_test.py
kill $TARGET_PID
```
**Scripts-based tests** (from `test-wheel-ubuntu` job):
```bash
# Tensor API perf test
export MOONCAKE_MASTER="127.0.0.1:50051"
export MOONCAKE_TE_META_DATA_SERVER="http://127.0.0.1:8080/metadata"
export MOONCAKE_PROTOCOL="tcp"
export LOCAL_HOSTNAME="127.0.0.1"
python scripts/test_tensor_api.py -n 1
python scripts/test_async_store.py
python scripts/test_copy_move_api.py
```
### Layer 3 — Static Checks (no services needed)
**CI equivalent:** `clang-format` and `spell-check` jobs
**Code format (changed files vs main):**
```bash
./scripts/code_format.sh --check --base origin/main
# Auto-fix:
./scripts/code_format.sh --base origin/main
```
**Spell check:**
```bash
# Requires typos tool: cargo install typos-cli
typos
```
**Pre-commit (runs all hooks):**
```bash
pip install pre-commit
pre-commit run --all-files
# Or just on staged files:
pre-commit run
```
## Build Configurations (from CI)
If the user needs to build first, here are the CI-equivalent cmake flags:
**Standard build with coverage (mirrors `build` job):**
```bash
mkdir build && cd build
cmake -G Ninja .. -DUSE_HTTP=ON -DUSE_CXL=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON -DENABLE_ASAN=ON -DCMAKE_BUILD_TYPE=Debug
cmake --build .
sudo cmake --install .
```
**All features ON (mirrors `build-flags` job):**
```bash
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
cmake --build .
sudo cmake --install .
```
**Transfer engine only:**
```bash
cd mooncake-transfer-engine
mkdir build && cd build
cmake -G Ninja .. -DUSE_ETCD=OFF -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON
cmake --build .
```
## Workflow: Diagnosing and Running Tests
### Step 1 — Understand what the user wants
Ask (or infer from context):
- All tests, or a specific subset?
- Did a specific CI job fail? Which one?
- Is the build already done, or do they need to build first?
### Step 2 — Check and Fix Prerequisites
**One-command setup** — this script checks all prerequisites and auto-fixes issues:
```bash
bash .claude/skills/mooncake-ci-local/scripts/check-prerequisites.sh
```
**What it checks:**
1. ✓ Build directory exists
2. ✓ mooncake package installed (auto-installs via cmake --install if missing)
3. ✓ ctest available
4. ✓ Restarts all services (etcd, metadata server) in clean state
5. ✓ Verifies all services are healthy
**If you need to build first:**
```bash
mkdir build && cd build
cmake -G Ninja .. -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CXL=ON -DSTORE_USE_ETCD=ON -DCMAKE_BUILD_TYPE=Debug
cmake --build .
sudo cmake --install .
```
**If script fails:**
- Build issues: See "Build Configurations" section below
- mooncake install fails: Try `pip install mooncake-wheel/dist/*.whl` manually
- etcd install fails: Download from https://github.com/etcd-io/etcd/releases
### Step 3 — Run and report
Run the relevant test layer. On failure:
1. Show the exact error message
2. Check if it's a service/env issue (most common) vs a real test failure
3. Suggest the fix (see common issues below)
## Common Local Test Issues
**"mooncake_master found in /usr/local/bin" error in run_tests.sh:**
The test expects mooncake_master to come from the Python package, not a system install.
```bash
# Remove the system-installed binary:
sudo rm /usr/local/bin/mooncake_master
# Or use the wheel-installed one:
pip install mooncake-wheel/dist/*.whl
```
**etcd port conflict:**
```bash
pkill etcd && sleep 1
etcd --advertise-client-urls http://127.0.0.1:2379 --listen-client-urls http://127.0.0.1:2379 &
```
**Metadata server port conflict:**
```bash
pkill -f bootstrap_server.py
pkill -f mooncake_http_metadata_server
```
**Tests hang (master not responding):**
```bash
pkill mooncake_master
sleep 2
mooncake_master --default_kv_lease_ttl=500 &
sleep 1
```
**torch/numpy not installed for tensor tests:**
```bash
pip install torch numpy safetensors packaging
```
**ctest shows no tests found:**
```bash
# Rebuild with unit tests enabled:
cd build
cmake .. -DBUILD_UNIT_TESTS=ON
cmake --build .
```
## Quick One-Liners
```bash
# Run ALL C++ tests (after building with etcd + metadata server running):
# Note: full suite takes 5-15 minutes depending on hardware
cd build && MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 ctest -j --output-on-failure
# Run only fast tests (skip slow integration tests):
cd build && MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 ctest -j --output-on-failure --exclude-regex "etcd|ha_test|redis"
# Run ALL Python tests:
mooncake_http_metadata_server --port 8080 & sleep 1 && cd mooncake-wheel/tests && MC_METADATA_SERVER=http://127.0.0.1:8080/metadata MC_FORCE_TCP=true bash ../../scripts/run_tests.sh
# Check code format (changed files only):
./scripts/code_format.sh --check --base origin/main
# Full pre-commit check:
pre-commit run --all-files
```

View File

@ -1,101 +0,0 @@
#!/bin/bash
# Mooncake CI Local Test Prerequisites Check
# Usage: bash check-prerequisites.sh
# This script checks and auto-fixes all prerequisites for running Mooncake CI tests locally.
set -e
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'
echo "🔍 Checking Mooncake CI test prerequisites..."
# 1. Check build directory
if [ ! -f build/CMakeCache.txt ]; then
echo -e "${RED}✗ Build directory not found or not built${NC}"
echo " → Run: mkdir build && cd build && cmake .. && cmake --build ."
exit 1
fi
echo -e "${GREEN}✓ Build exists${NC}"
# 2. Check mooncake installation
if ! python -c "import mooncake" 2>/dev/null; then
echo -e "${RED}✗ mooncake package not installed${NC}"
echo " → Fixing: Installing mooncake package..."
cd build && sudo cmake --install . && cd - >/dev/null
if ! python -c "import mooncake" 2>/dev/null; then
echo " → Alternative: pip install mooncake-wheel/dist/*.whl"
exit 1
fi
echo -e "${GREEN}✓ mooncake package installed${NC}"
else
echo -e "${GREEN}✓ mooncake package already installed${NC}"
fi
# 3. Check ctest availability
if ! command -v ctest &> /dev/null; then
echo -e "${RED}✗ ctest not found${NC}"
exit 1
fi
echo -e "${GREEN}✓ ctest available${NC}"
# 4. Kill and restart services (safest approach for local testing)
echo -e "\n${YELLOW}Cleaning up and restarting services...${NC}"
pkill -f "^etcd" || true
pkill -f bootstrap_server.py || true
pkill -f mooncake_http_metadata_server || true
sleep 1
# 5. Start etcd
if ! command -v etcd &> /dev/null; then
echo -e "${YELLOW}⚠ etcd not found, installing...${NC}"
ETCD_VER=v3.6.1
OS=$(uname -s | tr '[:upper:]' '[:lower:]')
ARCH=$(uname -m)
[ "$ARCH" = "x86_64" ] && ARCH="amd64"
DOWNLOAD_URL="https://github.com/etcd-io/etcd/releases/download/${ETCD_VER}/etcd-${ETCD_VER}-${OS}-${ARCH}.tar.gz"
echo " Downloading from: $DOWNLOAD_URL"
cd /tmp
wget -q "$DOWNLOAD_URL" && tar xzf "etcd-${ETCD_VER}-${OS}-${ARCH}.tar.gz" && \
sudo mv "etcd-${ETCD_VER}-${OS}-${ARCH}"/etcd* /usr/local/bin/
cd - >/dev/null
echo -e "${GREEN}✓ etcd installed${NC}"
fi
etcd --advertise-client-urls http://127.0.0.1:2379 --listen-client-urls http://127.0.0.1:2379 >/dev/null 2>&1 &
ETCD_PID=$!
sleep 2
if ! etcdctl --endpoints=http://127.0.0.1:2379 endpoint health &>/dev/null; then
echo -e "${RED}✗ etcd failed to start${NC}"
kill $ETCD_PID 2>/dev/null || true
exit 1
fi
echo -e "${GREEN}✓ etcd running (PID: $ETCD_PID)${NC}"
# 6. Start HTTP metadata server
if [ -f "mooncake-transfer-engine/example/http-metadata-server-python/bootstrap_server.py" ]; then
cd mooncake-transfer-engine/example/http-metadata-server-python
pip install -q aiohttp 2>/dev/null || true
python ./bootstrap_server.py >/dev/null 2>&1 &
METADATA_PID=$!
cd - >/dev/null
sleep 1
if curl -s http://127.0.0.1:8080/metadata > /dev/null 2>&1; then
echo -e "${GREEN}✓ HTTP Metadata server running (PID: $METADATA_PID)${NC}"
else
echo -e "${RED}✗ HTTP Metadata server failed to start${NC}"
kill $METADATA_PID $ETCD_PID 2>/dev/null || true
exit 1
fi
else
echo -e "${YELLOW}⚠ Metadata server script not found, skipping${NC}"
fi
echo -e "\n${GREEN}✅ All prerequisites ready!${NC}"
echo "Service PIDs: etcd=$ETCD_PID"
[ -n "$METADATA_PID" ] && echo "Metadata server PID: $METADATA_PID"
echo -e "\n${YELLOW}To kill services:${NC}"
echo " pkill -f '^etcd'"
echo " pkill -f bootstrap_server"

View File

@ -1,366 +0,0 @@
---
name: mooncake-troubleshoot
description: Automatically diagnose Mooncake deployment and runtime issues. Checks services (mooncake_master, metadata server), RDMA devices, environment variables, connectivity, memory limits, and analyzes logs for common error patterns. Use when Mooncake deployment fails, services won't start, connections fail, or you encounter runtime errors like "Error from etcd client", "No matched device found", "Failed to register memory", "NO_AVAILABLE_HANDLE", or any RDMA/networking issues. Also use when user asks to troubleshoot, debug, diagnose, or fix Mooncake problems.
---
# Mooncake Deployment Troubleshooting
You are a Mooncake deployment troubleshooting specialist. Your job is to systematically diagnose issues and provide actionable solutions based on the comprehensive troubleshooting knowledge from Mooncake documentation.
## Diagnostic Strategy
Run checks systematically, reporting findings as you go. Start with simple checks (services, connectivity) before diving into complex issues (RDMA, memory registration).
### 1. Service Status Check
Check if critical services are running:
```bash
# Check mooncake_master
ps aux | grep mooncake_master | grep -v grep
# Check port usage
netstat -tuln | grep -E '(50051|8080|2379|9003)'
# If using etcd
ps aux | grep etcd | grep -v grep
```
**Common issues:**
- `bind address already in use` → Port conflict, use different port with `--rpc_port`
- Master not running → Check startup logs for errors
### 2. Metadata Server Connectivity
The metadata server is critical for node discovery and coordination.
```bash
# Test etcd connectivity
curl -s http://127.0.0.1:2379/version
# Or test custom metadata server
curl -s $MC_METADATA_SERVER
# Check for proxy interference
echo "http_proxy: $http_proxy"
echo "https_proxy: $https_proxy"
```
**Common issues:**
- `Error from etcd client` → Metadata server unreachable
- **Fix:** Ensure etcd is bound to `0.0.0.0` not `127.0.0.1`:
```bash
etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://<your_ip>:2379
```
- **Fix:** Disable HTTP proxy:
```bash
unset http_proxy https_proxy
```
### 3. Environment Variables Check
Verify critical environment variables are set correctly:
```bash
# Display all MC_* variables
env | grep ^MC_
# Key variables to check:
echo "MC_METADATA_SERVER: $MC_METADATA_SERVER"
echo "MC_FORCE_TCP: $MC_FORCE_TCP"
echo "MC_LOG_LEVEL: $MC_LOG_LEVEL"
echo "MC_YLT_LOG_LEVEL: $MC_YLT_LOG_LEVEL"
echo "MC_MS_AUTO_DISC: $MC_MS_AUTO_DISC"
echo "MC_MS_FILTERS: $MC_MS_FILTERS"
echo "MC_GID_INDEX: $MC_GID_INDEX"
echo "MC_MTU: $MC_MTU"
echo "MC_IB_PORT: $MC_IB_PORT"
echo "MC_ENABLE_DEST_DEVICE_AFFINITY: $MC_ENABLE_DEST_DEVICE_AFFINITY"
```
**Key variables:**
- `MC_METADATA_SERVER` - Metadata server URL (required)
- `MC_FORCE_TCP=true` - Force TCP for testing without RDMA
- `MC_LOG_LEVEL=0` - Enable verbose logging (0=INFO, 1=WARNING, 2=ERROR)
- `MC_YLT_LOG_LEVEL=debug` - yalantinglibs log level
- `MC_MS_AUTO_DISC=1` - Enable topology auto-discovery (default)
- `MC_MS_FILTERS` - Filter specific RDMA devices (e.g., "mlx5_1,mlx5_2")
- `MC_GID_INDEX` - RDMA GID index (set if GID is all zeros)
- `MC_MTU` - RDMA MTU size
- `MC_ENABLE_DEST_DEVICE_AFFINITY=1` - Reduce QP creation (fix "Failed to create QP")
### 4. RDMA Device Check
Only run if RDMA is being used (skip if `MC_FORCE_TCP=true`):
```bash
# List RDMA devices
ibv_devices
# Check device details and status
ibv_devinfo
# Check for ACTIVE ports
ibv_devinfo | grep -A 10 "state:"
# Check GID addresses (should NOT be all zeros)
ibv_devinfo | grep -A 20 "GID"
# Check peer memory modules
lsmod | grep peer_mem
lsmod | grep nvidia_peer_mem
# Check QP count (if "Failed to create QP" error)
rdma resource show qp
```
**Common issues:**
- `No matched device found` → RDMA device name in config doesn't exist
- **Fix:** Use `ibv_devices` to get correct device names
- `Device XXX port not active` → RDMA port not in ACTIVE state
- **Fix:** Check cable connections, verify with `ibv_devinfo | grep state`
- **Fix:** Try different port with `MC_IB_PORT` environment variable
- GID all zeros → Wrong GID index
- **Fix:** Set `MC_GID_INDEX=1` (or 2, 3 depending on network)
- `Failed to create QP: Cannot allocate memory` → Too many QPs created
- **Fix:** Set `MC_ENABLE_DEST_DEVICE_AFFINITY=1`
### 5. Memory and Resource Limits
Check system limits that affect RDMA memory registration:
```bash
# Check ulimits
ulimit -a
# Focus on max locked memory
ulimit -l
# Check RDMA device memory limits
ibv_devinfo -v | grep max_mr_size
# Check dmesg for memory errors
dmesg -T | tail -50 | grep -i "out of mr size"
```
**Common issues:**
- `Failed to register memory: Input/output error` → Memory registration limit exceeded
- **Diagnostic:** Check `max_mr_size` with `ibv_devinfo -v`
- **Fix:** Reduce memory allocation or split into smaller chunks
- Cannot allocate memory → ulimit restriction
- **Fix:** Set unlimited locked memory:
```bash
ulimit -l unlimited
```
- **Permanent fix:** Add to `/etc/security/limits.conf`:
```
* soft memlock unlimited
* hard memlock unlimited
```
### 6. Network Connectivity
Test connectivity between nodes:
```bash
# Test basic RDMA connectivity
ib_write_bw -d <device_name> -R
# On peer node:
ib_write_bw -d <device_name> -R <server_ip>
# Test GPU Direct RDMA (if CUDA enabled)
ib_write_bw -d <device_name> -R -x gdr
# On peer node:
ib_write_bw -d <device_name> -R -x gdr <server_ip>
# Test DNS resolution
nslookup <connectable_name>
ping <connectable_name>
```
**Common issues:**
- `connection refused` → Incorrect `connectable_name` or `rpc_port`
- **Fix:** Ensure `connectable_name` is NOT loopback (127.0.0.1/localhost)
- **Fix:** Use actual LAN/WAN IP or valid hostname
- `Failed to exchange handshake` → RDMA connection setup failure
- **Fix:** Verify MTU matches: set `MC_MTU` environment variable
- **Fix:** Verify GID is valid (not all zeros)
- **Fix:** Test with `ib_send_bw` between nodes first
### 7. Log Analysis
Search logs for common error patterns and their meanings:
**Metadata/Connectivity Errors:**
- `Error from etcd client` → Cannot connect to metadata server
- `ERR_METADATA` → Metadata server communication failed
- `ERR_DNS` → Invalid `local_server_name` (not valid DNS/IP)
**RDMA Errors:**
- `No matched device found` → RDMA device name doesn't exist
- `Device XXX port not active` → RDMA port not in ACTIVE state
- `Failed to exchange handshake description` → RDMA handshake failed
- `Failed to modify QP to RTR, check mtu, gid, peer lid, peer qp num` → MTU/GID mismatch
- `Failed to register memory` → Memory registration limit exceeded
- `Failed to create QP` → Too many QPs, enable `MC_ENABLE_DEST_DEVICE_AFFINITY=1`
- `Worker: Process failed for slice` → Network instability
- `work request flushed error` → Cascading error (find first error)
**Store Errors:**
- `NO_AVAILABLE_HANDLE` (-200) → Memory pool exhausted
- **Fix:** Increase `global_segment_size` in setup
- **Fix:** Check eviction is working (look for eviction logs)
- `LEASE_EXPIRED` (-707) → Lease expired during transfer
- **Fix:** Increase `default_kv_lease_ttl` in master startup
- `OBJECT_NOT_FOUND` (-704) → Object doesn't exist
- `SEGMENT_NOT_FOUND` (-101) → No available segments
- `Failed to get description of XXX` → Segment name mismatch
- **Fix:** Ensure segment name matches `local_hostname` from peer
**Port/Service Errors:**
- `bind address already in use` → Port conflict
- **Fix:** Use different port: `--rpc_port=50052`
### 8. Configuration Validation
Verify configuration is correct:
```bash
# Check connectable_name is not loopback
hostname -I
# Verify master startup flags
ps aux | grep mooncake_master
# Check if using correct protocol
env | grep MC_FORCE_TCP
```
**Critical checks:**
- `connectable_name` must be non-loopback IP or valid hostname
- MTU and GID configurations must match network environment
- RDMA device names must exist on the machine
- Ports must not be in use by other services
## Error Code Quick Reference
### Transfer Engine Error Codes
| Code | Name | Meaning | Fix |
|------|------|---------|-----|
| 0 | Success | Normal execution | - |
| -12 | ERR_ADDRESS_NOT_REGISTERED | Memory not registered | Register memory before use |
| -14 | ERR_DEVICE_NOT_FOUND | RDMA device not found | Check device name with `ibv_devices` |
| -16 | ERR_DNS | Invalid local_server_name | Use valid IP/hostname |
| -19 | ERR_REJECT_HANDSHAKE | Peer rejected handshake | Check peer logs for reason |
| -20 | ERR_METADATA | Metadata server unreachable | Check etcd/HTTP server |
### Store Error Codes
| Code | Name | Meaning | Fix |
|------|------|---------|-----|
| 0 | Success | Operation successful | - |
| -200 | NO_AVAILABLE_HANDLE | Memory pool exhausted | Increase segment size |
| -707 | LEASE_EXPIRED | Lease expired | Increase lease TTL |
| -704 | OBJECT_NOT_FOUND | Object doesn't exist | Check object key |
| -101 | SEGMENT_NOT_FOUND | No available segments | Check segment registration |
| -900 | RPC_FAIL | RPC failed | Check network/master |
| -1000 | ETCD_OPERATION_ERROR | etcd operation failed | Check etcd status |
## Output Format
Provide a structured diagnostic report:
```
🔍 MOONCAKE DEPLOYMENT DIAGNOSTICS
==================================
✅ PASSED CHECKS:
- Service status: mooncake_master running on port 50051
- Metadata server: etcd accessible at http://127.0.0.1:2379
- Environment: MC_METADATA_SERVER set correctly
- [other passing checks]
❌ FAILED CHECKS:
- RDMA device: mlx5_0 port not ACTIVE (state: PORT_DOWN)
- Memory limits: max locked memory is 64KB (too low)
- [other failures with specific error messages]
⚠️ WARNINGS:
- GID index not set, may cause connection issues
- HTTP proxy variables set, may interfere with metadata server
- [other potential issues]
🔧 RECOMMENDED FIXES:
1. Fix RDMA port status:
- Check physical cable connections
- Verify driver configuration
- Command: ibv_devinfo | grep -A 10 "state:"
2. Increase memory limits:
ulimit -l unlimited
# Or permanently in /etc/security/limits.conf:
* soft memlock unlimited
* hard memlock unlimited
3. Set GID index:
export MC_GID_INDEX=1
4. Disable HTTP proxy:
unset http_proxy https_proxy
📋 SUMMARY:
[Brief 2-3 sentence conclusion about deployment health and next steps]
```
## Troubleshooting Workflow
1. **Start simple**: Check services and basic connectivity first
2. **Read logs carefully**: First error is usually root cause (subsequent errors cascade)
3. **Test incrementally**: Use `MC_FORCE_TCP=true` to isolate RDMA issues
4. **Verify basics**: Check connectable_name, ports, env vars before deep diving
5. **Use diagnostic tools**: ibv_devices, ibv_devinfo, ib_write_bw, curl
6. **Reference documentation**: Check error codes and troubleshooting guide
## Quick Fix Commands
**Start metadata server properly:**
```bash
etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://<your_ip>:2379
```
**Enable verbose logging:**
```bash
export MC_LOG_LEVEL=0
export MC_YLT_LOG_LEVEL=debug
```
**Force TCP mode for testing:**
```bash
export MC_FORCE_TCP=true
```
**Fix memory limits:**
```bash
ulimit -l unlimited
```
**Fix too many QPs:**
```bash
export MC_ENABLE_DEST_DEVICE_AFFINITY=1
```
**Fix GID issues:**
```bash
export MC_GID_INDEX=1 # or 2, 3 depending on network
```
**Use different port:**
```bash
mooncake_master --rpc_port=50052
```
Now execute the diagnostic checks systematically and provide the structured report.

View File

@ -23,20 +23,11 @@ RUN apt-get install -y libibverbs-dev \
libhiredis-dev \
libyaml-cpp-dev \
libjemalloc-dev \
libzstd-dev \
libmsgpack-dev \
libgflags-dev \
pkg-config \
patchelf
RUN GO_VERSION="1.23.8" && \
ARCH=$(uname -m) && \
if [ "$ARCH" = "aarch64" ]; then GOARCH="arm64"; \
elif [ "$ARCH" = "x86_64" ]; then GOARCH="amd64"; \
else echo "Unsupported architecture: $ARCH" && exit 1; fi && \
wget https://go.dev/dl/go${GO_VERSION}.linux-${GOARCH}.tar.gz \
&& tar -C /usr/local -xzf go${GO_VERSION}.linux-${GOARCH}.tar.gz \
&& rm go${GO_VERSION}.linux-${GOARCH}.tar.gz
RUN wget https://go.dev/dl/go1.22.12.linux-amd64.tar.gz \
&& tar -C /usr/local -xzf go1.22.12.linux-amd64.tar.gz
RUN git clone https://github.com/alibaba/yalantinglibs.git \
&& cd yalantinglibs \

10
.github/CODEOWNERS vendored
View File

@ -8,16 +8,14 @@
.github @stmatengss @ykwd @Ann-1024 @luketong777
/docs @ShangmingCai @stmatengss @ykwd
/mooncake-ep @UNIDY2002 @ympcMark @yuechen-sys
/mooncake-integration/transfer_engine @ShangmingCai @alogfans
/mooncake-ep @UNIDY2002 @ympcMark
/mooncake-integration/ep @UNIDY2002 @ympcMark
/mooncake-integration/transfer_engine @ShangmingCai @alogfans
/mooncake-integration/store @ykwd @stmatengss
/mooncake-pg @UNIDY2002 @ympcMark @yuechen-sys
/mooncake-pg @UNIDY2002 @ympcMark
/mooncake-store @ykwd @stmatengss @XucSh @YiXR
/mooncake-store/*/ha/ @Libotry @YiXR @00fish0
/mooncake-transfer-engine @alogfans @doujiang24 @chestnut-Q
/mooncake-transfer-engine/*/transport/hip_transport/ @alogfans @amd-arozanov
/mooncake-transfer-engine/*/transport/ascend_transport/ @alogfans @ascend-direct-dev
/mooncake-transfer-engine/*/transport/efa_transport/ @alogfans @whn09
/mooncake-wheel @ShangmingCai @stmatengss
/scripts/tone_tests @luketong777
/scripts/ascend/ @ascend-direct-dev @VNightMare @MingYang119

View File

@ -6,25 +6,14 @@ on:
pull_request:
branches: [ "main" ]
types: [opened, synchronize, reopened, labeled]
workflow_dispatch: {}
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.ref_name }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: true
jobs:
build:
needs: [spell-check, clang-format, check-paths]
if: >-
(needs.check-paths.outputs.should-run-downstream == 'true' ||
github.event_name == 'workflow_dispatch') &&
(github.event_name == 'push' ||
github.event_name == 'workflow_dispatch' ||
github.event.action == 'opened' ||
contains(github.event.pull_request.labels.*.name, 'run-ci'))
github.event_name == 'push' ||
github.event.action == 'opened' ||
contains(github.event.pull_request.labels.*.name, 'run-ci')
runs-on: ubuntu-22.04
strategy:
matrix:
@ -35,8 +24,6 @@ jobs:
steps:
- uses: actions/checkout@v4
with:
persist-credentials: false
- name: Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@v5
@ -67,10 +54,10 @@ jobs:
method: 'network'
sub-packages: '["nvcc"]'
- name: Install coverage tools and build utilities
- name: Install coverage tools
run: |
sudo apt-get update
sudo apt-get install -y lcov gcovr ninja-build
sudo apt-get install -y lcov gcovr
- name: Set up coverage compilation flags
run: |
@ -100,14 +87,14 @@ 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 .. -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
run: |
cd build
cmake --build .
sudo cmake --install .
make -j4
sudo make install
shell: bash
- name: Build nvlink_allocator.so
@ -125,42 +112,12 @@ jobs:
python ./bootstrap_server.py &
shell: bash
- name: Run Go store binding integration tests
run: |
$GITHUB_WORKSPACE/build/mooncake-store/src/mooncake_master \
--eviction_high_watermark_ratio=0.95 \
--cluster_id=ci_go_test_cluster \
--port 50051 &
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 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
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
- name: Test (in build env) with coverage
run: |
cd build
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
ldconfig -v || echo "always continue"
MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 ctest -j --output-on-failure
shell: bash
- name: Drain HTTP E2E test
if: matrix.python-version == '3.12'
run: |
cd build
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
# Keep the sanitizer gate on the C++ integration test. The Python
# drain script is manual/nightly only because pybind + ASan teardown in
# a Python host process is not stable.
DEFAULT_KV_LEASE_TTL=500 ./mooncake-store/tests/task_integration_test --gtest_filter='TaskExecutorIntegrationTest.DrainJobCompleteFlow'
MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 make test -j ARGS="-V"
shell: bash
- name: Generate coverage report
@ -217,6 +174,13 @@ jobs:
echo "✅ Coverage collected successfully"
fi
- name: Test mooncake conductor
run: |
cd mooncake-conductor/conductor-ctrl
go mod tidy
go test ./... -v -cover
shell: bash
- name: Generate Python version tag
id: generate_tag_build
run: |
@ -236,20 +200,14 @@ jobs:
path: mooncake-wheel/dist-py${{ steps.generate_tag_build.outputs.python_version_tag }}/*.whl
build-musa:
needs: [spell-check, clang-format, check-paths]
if: >-
(needs.check-paths.outputs.should-run-downstream == 'true' ||
github.event_name == 'workflow_dispatch') &&
(github.event_name == 'push' ||
github.event_name == 'workflow_dispatch' ||
github.event.action == 'opened' ||
contains(github.event.pull_request.labels.*.name, 'run-ci'))
github.event_name == 'push' ||
github.event.action == 'opened' ||
contains(github.event.pull_request.labels.*.name, 'run-ci')
runs-on: ubuntu-22.04
container: mthreads/musa:rc4.3.0-devel-ubuntu22.04-amd64
steps:
- uses: actions/checkout@v4
with:
persist-credentials: false
- name: Mark repository as safe
run: git config --global --add safe.directory $GITHUB_WORKSPACE
@ -258,29 +216,22 @@ jobs:
- name: Configure project
run: |
apt update -y
apt install -y ninja-build
bash -x dependencies.sh -y
mkdir build
cd build
cmake -G Ninja .. -DUSE_MUSA=ON -DUSE_MNNVL=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON -DUSE_CXL=ON -DUSE_TCP=ON -DBUILD_UNIT_TESTS=OFF -DBUILD_EXAMPLES=OFF
cmake .. -DUSE_MUSA=ON -DUSE_MNNVL=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON -DUSE_CXL=ON -DUSE_TCP=ON -DBUILD_UNIT_TESTS=OFF -DBUILD_EXAMPLES=OFF
shell: bash
- name: Build project
run: |
cd build
source ~/.bashrc
cmake --build .
cmake --install .
make -j
make install
shell: bash
test-wheel-ubuntu:
needs: [spell-check, clang-format, build-flags]
if: >-
needs.build-flags.result == 'success' &&
(github.event_name == 'push' ||
github.event_name == 'workflow_dispatch' ||
github.event.action == 'opened' ||
contains(github.event.pull_request.labels.*.name, 'run-ci'))
needs: build-flags
strategy:
matrix:
ubuntu-version: [ubuntu-22.04, ubuntu-24.04]
@ -288,8 +239,6 @@ jobs:
runs-on: ${{ matrix.ubuntu-version }}
steps:
- uses: actions/checkout@v4
with:
persist-credentials: false
- name: Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@v5
@ -345,14 +294,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
@ -403,18 +344,6 @@ jobs:
python scripts/test_copy_move_api.py
shell: bash
- name: Run Python Drain HTTP E2E Test (CI check)
env:
MOONCAKE_MASTER: "127.0.0.1:50051"
MOONCAKE_TE_META_DATA_SERVER: "http://127.0.0.1:8080/metadata"
MOONCAKE_PROTOCOL: "tcp"
LOCAL_HOSTNAME: "127.0.0.1"
run: |
source test_env/bin/activate
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
python scripts/test_drain_http_api.py --timeout-sec 90
shell: bash
- name: Run RPC Communicator Bandwidth Test
run: |
source test_env/bin/activate
@ -425,12 +354,15 @@ jobs:
kill $SERVER_PID 2>/dev/null || true
wait $SERVER_PID 2>/dev/null || true
- name: Test Mooncake PyTorch Backend (CPU Only)
- name: Test Mooncake EP Backend (CPU Only)
env:
MC_FORCE_TCP: "true"
run: |
source test_env/bin/activate
python -m unittest mooncake-wheel.tests.test_mooncake_backend_cpu
# Disable these tests in CI as they fail occasionally.
# python -m unittest mooncake-wheel.tests.test_mooncake_backend_elastic
# python -m unittest mooncake-wheel.tests.test_mooncake_backend_p2p_cpu
shell: bash
- name: Test Safetensor Functions
@ -441,14 +373,10 @@ jobs:
shell: bash
build-flags:
needs: [spell-check, clang-format, check-paths]
if: >-
(needs.check-paths.outputs.should-run-downstream == 'true' ||
github.event_name == 'workflow_dispatch') &&
(github.event_name == 'push' ||
github.event_name == 'workflow_dispatch' ||
github.event.action == 'opened' ||
contains(github.event.pull_request.labels.*.name, 'run-ci'))
github.event_name == 'push' ||
github.event.action == 'opened' ||
contains(github.event.pull_request.labels.*.name, 'run-ci')
runs-on: ubuntu-22.04
strategy:
matrix:
@ -456,13 +384,12 @@ jobs:
env:
CI: "true"
BUILD_WITH_EP: "1"
EP_TORCH_VERSIONS: "2.9.0;2.9.1;2.10.0"
TORCH_CUDA_ARCH_LIST: "8.0;9.0"
SCCACHE_GHA_ENABLED: "true"
steps:
- uses: actions/checkout@v4
with:
persist-credentials: false
- name: Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@v5
@ -503,14 +430,10 @@ jobs:
- name: Install dependencies
run: |
sudo apt update -y
sudo apt install -y ninja-build
sudo bash -x dependencies.sh -y
df -h
shell: bash
- name: Install Rust toolchain
uses: dtolnay/rust-toolchain@stable
- name: Build transfer engine only
run: |
cd mooncake-transfer-engine
@ -518,9 +441,9 @@ 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 --build .
sudo cmake --install .
cmake .. -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"
make -j4
sudo make install
df -h
shell: bash
@ -528,7 +451,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 .. -DUSE_ETCD=ON -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_STORE=ON -DWITH_P2P_STORE=ON -DWITH_EP=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
@ -537,39 +460,32 @@ jobs:
export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
cd build
cmake --build .
sudo cmake --install .
make -j4
sudo make install
df -h
shell: bash
- name: Configure project with unit tests and examples
run: |
cd build
cmake -G Ninja .. -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DWITH_STORE_RUST=ON -DENABLE_SCCACHE=ON
cmake .. -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON
shell: bash
# TODO: lack WITH_RUST_EXAMPLE
- name: Build project with unit tests and examples
run: |
export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
cd build
cmake --build .
sudo cmake --install .
shell: bash
- name: Check Mooncake Store Rust bindings and example
run: |
cd mooncake-store/rust
MOONCAKE_STORE_LIB_DIR=$GITHUB_WORKSPACE/build/mooncake-store/src \
MOONCAKE_STORE_INCLUDE_DIR=$GITHUB_WORKSPACE/mooncake-store/include \
cargo check --example basic_usage --tests
make -j4
sudo make install
shell: bash
- name: Configure project
run: |
cd build
rm -r */tests
cmake -G Ninja .. -DBUILD_UNIT_TESTS=OFF -DBUILD_EXAMPLES=OFF -DUSE_HTTP=ON -DENABLE_SCCACHE=ON -DUSE_CXL=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.9.0;2.9.1;2.10.0;2.11.0"
cmake .. -DBUILD_UNIT_TESTS=OFF -DBUILD_EXAMPLES=OFF -DUSE_HTTP=ON -DENABLE_SCCACHE=ON -DUSE_CXL=ON
shell: bash
- name: Build project
@ -577,8 +493,8 @@ jobs:
export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
cd build
cmake --build .
sudo cmake --install .
make -j4
sudo make install
shell: bash
- name: Build nvlink_allocator.so
@ -611,19 +527,13 @@ jobs:
build-docker:
name: Build Docker Image
needs: [spell-check, clang-format, check-paths]
if: >-
(needs.check-paths.outputs.should-run-downstream == 'true' ||
github.event_name == 'workflow_dispatch') &&
(github.event_name == 'push' ||
github.event_name == 'workflow_dispatch' ||
github.event.action == 'opened' ||
contains(github.event.pull_request.labels.*.name, 'run-ci'))
github.event_name == 'push' ||
github.event.action == 'opened' ||
contains(github.event.pull_request.labels.*.name, 'run-ci')
runs-on: ubuntu-22.04
steps:
- uses: actions/checkout@v4
with:
persist-credentials: false
- name: Set up Docker Buildx
uses: docker/setup-buildx-action@v2
@ -639,15 +549,12 @@ jobs:
name: Spell Check with Typos
if: >-
github.event_name == 'push' ||
github.event_name == 'workflow_dispatch' ||
github.event.action == 'opened' ||
contains(github.event.pull_request.labels.*.name, 'run-ci')
runs-on: ubuntu-22.04
steps:
- name: Checkout Actions Repository
uses: actions/checkout@v4
with:
persist-credentials: false
- name: Spell Check Repo
uses: crate-ci/typos@v1.30.2
@ -655,7 +562,6 @@ jobs:
name: Check code format
if: >-
github.event_name == 'push' ||
github.event_name == 'workflow_dispatch' ||
github.event.action == 'opened' ||
contains(github.event.pull_request.labels.*.name, 'run-ci')
runs-on: ubuntu-22.04
@ -664,7 +570,6 @@ jobs:
uses: actions/checkout@v4
with:
fetch-depth: 0 # Need full history for branch comparison
persist-credentials: false
- name: Install clang-format 20
run: |
@ -699,92 +604,3 @@ jobs:
echo "Comparing against: ${BASE_REF}"
./scripts/code_format.sh --check --base "${BASE_REF}"
shell: bash
check-paths:
if: >-
github.event_name == 'push' ||
github.event_name == 'workflow_dispatch' ||
github.event.action == 'opened' ||
contains(github.event.pull_request.labels.*.name, 'run-ci')
runs-on: ubuntu-latest
outputs:
should-run-downstream: ${{ steps.dispatch-override.outputs.src || steps.filter.outputs.src }}
steps:
# workflow_dispatch has no PR/push diff context — skip paths-filter and default to true
- name: Default to true for workflow_dispatch
id: dispatch-override
if: github.event_name == 'workflow_dispatch'
run: echo "src=true" >> $GITHUB_OUTPUT
- uses: actions/checkout@v4
if: github.event_name != 'workflow_dispatch'
with:
fetch-depth: 2
persist-credentials: false
- uses: dorny/paths-filter@v3
if: github.event_name != 'workflow_dispatch'
id: filter
with:
filters: |
src:
- 'mooncake-*/**'
- 'extern/**'
- 'CMakeLists.txt'
- 'dependencies.sh'
- 'scripts/**'
- '.github/workflows/**'
build-wheel-cu13:
needs: [spell-check, clang-format, check-paths]
if: >-
(needs.check-paths.outputs.should-run-downstream == 'true' ||
github.event_name == 'workflow_dispatch') &&
(github.event_name == 'push' ||
github.event_name == 'workflow_dispatch' ||
github.event.action == 'opened' ||
contains(github.event.pull_request.labels.*.name, 'run-ci'))
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()
needs:
- spell-check
- clang-format
- build
- build-musa
- build-flags
- build-docker
- test-wheel-ubuntu
- build-wheel-cu13
- ascend-test
- integration-test
runs-on: ubuntu-latest
steps:
- name: Check required job results
run: |
failing=$(echo "$NEEDS_JSON" | jq -r '
to_entries[] |
select(.value.result != "success" and .value.result != "skipped") |
"\(.key): \(.value.result)"')
if [ -n "$failing" ]; then
echo "::error::The following jobs failed or were cancelled:"
echo "$failing"
exit 1
fi
echo "All checks passed or were acceptably skipped."
env:
NEEDS_JSON: ${{ toJSON(needs) }}

View File

@ -1,364 +0,0 @@
name: 'CI Test on ASCEND Platform'
on:
workflow_call:
inputs:
checkout_ref:
description: 'Git ref to checkout (PR head SHA for pull_request_target)'
required: false
type: string
jobs:
build-and-test:
if: github.repository == 'kvcache-ai/Mooncake'
runs-on: self-hosted
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
--device /dev/davinci_manager --device /dev/devmm_svm --device /dev/hisi_hdc --ulimit nproc=65535:65535
env:
GITHUB_ACTIONS: "true"
LD_PRELOAD: "/usr/lib64/libjemalloc.so.2:"
volumes:
- /usr/local/dcmi:/usr/local/dcmi
- /usr/local/Ascend/driver/:/usr/local/Ascend/driver/
- /etc/ascend_install.info:/etc/ascend_install.info
- /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/
else
echo "Error: ../pybind11 does not exist. Cannot copy pybind11."
exit 1
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
mkdir -p build
cd build
cmake .. \
-DUSE_ASCEND_DIRECT=ON \
-DBUILD_EXAMPLES=OFF \
-DBUILD_UNIT_TESTS=OFF
- name: Build
shell: bash
run: |
source /usr/local/Ascend/cann-9.0.0/set_env.sh
echo "Building..."
cd build
cmake --build . -j$(nproc)
cmake --install .
echo "Mooncake installed successfully."
- name: Run Hixl Mooncake Store Test
shell: bash
run: |
source /usr/local/Ascend/cann-9.0.0/set_env.sh
set -e
export ASCEND_PROCESS_LOG_PATH=/tmp/hixl-test-log/
export ASCEND_GLOBAL_LOG_LEVEL=3
echo "=== Cloning Hixl repository ==="
cd ..
rm -rf hixl
git clone https://gitcode.com/cann/hixl.git
cd hixl/examples/third_parties/mooncake_store/python/
export LD_LIBRARY_PATH=/usr/local/lib:${LD_LIBRARY_PATH}
echo "=== Starting Mooncake Master ==="
# Find mooncake_master binary
MOONCAKE_MASTER=$(find /usr/local/bin /usr/bin -name "mooncake_master" -type f 2>/dev/null | head -1)
if [ -z "$MOONCAKE_MASTER" ]; then
# Try finding in build directory
MOONCAKE_MASTER=$(find $GITHUB_WORKSPACE/build -name "mooncake_master" -type f 2>/dev/null | head -1)
fi
if [ -z "$MOONCAKE_MASTER" ]; then
echo "Error: mooncake_master binary not found"
exit 1
fi
echo "Found mooncake_master at: $MOONCAKE_MASTER"
# Start Mooncake master in background
$MOONCAKE_MASTER \
--enable_http_metadata_server=true \
--http_metadata_server_host=0.0.0.0 \
--http_metadata_server_port=8080 \
> /tmp/mooncake_master.log 2>&1 &
MASTER_PID=$!
echo "Mooncake Master started with PID: $MASTER_PID"
# Wait for master to be ready
echo "Waiting for Mooncake Master to initialize..."
sleep 5
# 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
exit 1
fi
echo "Mooncake Master is running"
echo "=== Running Hixl Mooncake Store Tests ==="
# List of test cases to run
TEST_CASES=(
"batch_put_get_sample.py"
"batch_put_get_multi_buffers_sample.py"
)
# List of test scenarios (HCCL_INTRA_ROCE_ENABLE settings)
TEST_SCENARIOS=(
"HCCL_INTRA_ROCE_ENABLE=1"
"HCCL_INTRA_ROCE_ENABLE_UNSET"
)
# Track test results
FAILED_TESTS=()
PASSED_TESTS=()
# Run each test scenario
for scenario in "${TEST_SCENARIOS[@]}"; do
echo ""
echo "========================================="
echo "Running scenario: $scenario"
echo "========================================="
# Configure environment variables for the current scenario
if [ "$scenario" = "HCCL_INTRA_ROCE_ENABLE=1" ]; then
export HCCL_INTRA_ROCE_ENABLE=1
unset ASCEND_BUFFER_POOL
echo "HCCL_INTRA_ROCE_ENABLE is set to 1, ASCEND_BUFFER_POOL is unset"
else
unset HCCL_INTRA_ROCE_ENABLE
export ASCEND_BUFFER_POOL=4:8
echo "HCCL_INTRA_ROCE_ENABLE is not set, ASCEND_BUFFER_POOL is set to 4:8"
fi
# Run each test case in the current scenario
for test_case in "${TEST_CASES[@]}"; do
echo ""
echo "-----------------------------------------"
echo "Test: $test_case"
echo "-----------------------------------------"
if [ ! -f "$test_case" ]; then
echo "Warning: Test file $test_case not found, skipping..."
continue
fi
# Run the test with 2 devices in distributed mode
# Run rank 0 on device 0
python3 $test_case \
--device_id=0 \
--rank=0 \
--world_size=2 \
--distributed \
2>&1 | tee "/tmp/hixl_test_${scenario//=/}_${test_case%.py}_rank0.log" &
PID0=$!
# Run rank 1 on device 1
python3 $test_case \
--device_id=2 \
--rank=1 \
--world_size=2 \
--distributed \
2>&1 | tee "/tmp/hixl_test_${scenario//=/}_${test_case%.py}_rank1.log" &
PID1=$!
# Wait for both processes to complete
wait $PID0
TEST_RESULT0=$?
wait $PID1
TEST_RESULT1=$?
# Check test results
if [ $TEST_RESULT0 -eq 0 ] && [ $TEST_RESULT1 -eq 0 ]; then
echo "✓ $test_case PASSED (scenario: $scenario)"
PASSED_TESTS+=("$scenario:$test_case")
else
echo "✗ $test_case FAILED (scenario: $scenario)"
if [ $TEST_RESULT0 -ne 0 ]; then
echo " Rank 0 failed with code: $TEST_RESULT0"
fi
if [ $TEST_RESULT1 -ne 0 ]; then
echo " Rank 1 failed with code: $TEST_RESULT1"
fi
FAILED_TESTS+=("$scenario:$test_case")
fi
done
done
echo ""
echo "========================================="
echo "Test Summary"
echo "========================================="
echo "Passed tests: ${#PASSED_TESTS[@]}"
for test in "${PASSED_TESTS[@]}"; do
echo " ✓ $test"
done
echo ""
echo "Failed tests: ${#FAILED_TESTS[@]}"
for test in "${FAILED_TESTS[@]}"; do
echo " ✗ $test"
done
# Cleanup: Stop Mooncake Master
echo ""
echo "Stopping Mooncake Master..."
kill $MASTER_PID 2>/dev/null || true
wait $MASTER_PID 2>/dev/null || true
# Exit with error if any tests failed
if [ ${#FAILED_TESTS[@]} -gt 0 ]; then
echo ""
echo "Some tests failed!"
exit 1
fi
echo ""
echo "All Hixl Mooncake Store tests completed successfully!"
- name: Test Summary
if: always()
shell: bash
run: |
echo "CI Test completed"
- name: Upload Test Logs
if: always()
uses: actions/upload-artifact@v4
with:
name: test-logs-${{ github.run_number }}
path: |
/tmp/hixl-test-log/*
retention-days: 30
if-no-files-found: warn

View File

@ -1,10 +1,18 @@
name: 'Build Wheel (CUDA 13)'
on:
workflow_call: {}
push:
branches: [ "main" ]
pull_request:
branches: [ "main" ]
types: [opened, synchronize, reopened, labeled]
jobs:
build-wheel-cu13:
if: >-
github.event_name == 'push' ||
github.event.action == 'opened' ||
contains(github.event.pull_request.labels.*.name, 'run-ci')
runs-on: ubuntu-22.04
strategy:
matrix:
@ -12,13 +20,12 @@ jobs:
env:
BUILD_WITH_EP: "1"
CU13_BUILD: "1"
EP_TORCH_VERSIONS: "2.9.0;2.9.1;2.10.0"
TORCH_CUDA_ARCH_LIST: "8.0;9.0"
SCCACHE_GHA_ENABLED: "true"
steps:
- uses: actions/checkout@v4
with:
persist-credentials: false
- name: Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@v5
@ -59,7 +66,6 @@ jobs:
- name: Install dependencies
run: |
sudo apt update -y
sudo apt install -y ninja-build
sudo bash -x dependencies.sh -y
df -h
shell: bash
@ -68,14 +74,13 @@ jobs:
run: |
mkdir build
cd build
cmake -G Ninja .. \
cmake .. \
-DUSE_ETCD=ON \
-DUSE_REDIS=ON \
-DUSE_HTTP=ON \
-DWITH_STORE=ON \
-DWITH_P2P_STORE=ON \
-DWITH_EP=ON \
-DEP_TORCH_VERSIONS="2.9.0;2.9.1;2.10.0;2.11.0" \
-DWITH_METRICS=ON \
-DBUILD_UNIT_TESTS=OFF \
-DBUILD_EXAMPLES=ON \
@ -91,8 +96,8 @@ jobs:
export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH
export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH
cd build
cmake --build .
sudo cmake --install .
make -j4
sudo make install
df -h
shell: bash

View File

@ -4,10 +4,6 @@ on:
# Runs on pushes targeting the default branch
push:
branches: ["main"]
paths:
- 'docs/**'
- 'requirements_docs.txt'
- '.github/workflows/deploy.yml'
# Allows you to run this workflow manually from the Actions tab
workflow_dispatch:
@ -35,9 +31,7 @@ jobs:
steps:
- name: Checkout
uses: actions/checkout@v4
with:
persist-credentials: false
- name: Setup Python
uses: actions/setup-python@v4
with:

View File

@ -1,94 +0,0 @@
name: E2E CI
on:
pull_request_target:
branches: ["main"]
types: [labeled]
workflow_dispatch:
inputs:
pr_number:
description: 'PR number that triggered this'
required: false
type: string
pr_sha:
description: 'PR head SHA to checkout'
required: false
type: string
triggered_by:
description: 'User who triggered this'
required: false
type: string
permissions:
contents: read
pull-requests: write
concurrency:
group: e2e-ci-${{ github.event.pull_request.number || inputs.pr_number || github.sha }}
cancel-in-progress: true
jobs:
ascend-test:
if: >
github.event_name == 'workflow_dispatch' ||
github.event.label.name == 'run-e2e-ci'
uses: ./.github/workflows/ci_ascend.yml
with:
checkout_ref: ${{ inputs.pr_sha || github.event.pull_request.head.sha }}
secrets: inherit
integration-test:
if: >
github.event_name == 'workflow_dispatch' ||
github.event.label.name == 'run-e2e-ci'
uses: ./.github/workflows/integration-test.yml
with:
pr_sha: ${{ inputs.pr_sha || github.event.pull_request.head.sha }}
pr_number: ${{ inputs.pr_number || github.event.pull_request.number }}
secrets: inherit
e2e-gate:
name: E2E Gate
if: >
always() &&
(github.event_name == 'workflow_dispatch' ||
github.event.label.name == 'run-e2e-ci')
needs:
- ascend-test
- integration-test
runs-on: ubuntu-latest
steps:
- name: Check E2E results
run: |
echo "PR: #${{ inputs.pr_number || github.event.pull_request.number }}"
echo "SHA: ${{ inputs.pr_sha || github.event.pull_request.head.sha }}"
failing=$(echo "$NEEDS_JSON" | jq -r '
to_entries[] |
select(.value.result != "success" and .value.result != "skipped") |
"\(.key): \(.value.result)"')
if [ -n "$failing" ]; then
echo "::error::The following E2E jobs failed:"
echo "$failing"
exit 1
fi
echo "All E2E checks passed."
env:
NEEDS_JSON: ${{ toJSON(needs) }}
cleanup-label:
name: Cleanup E2E Label
if: >
always() &&
github.event_name != 'workflow_dispatch' &&
github.event.label.name == 'run-e2e-ci'
needs:
- e2e-gate
runs-on: ubuntu-latest
steps:
- name: Remove run-e2e-ci label
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: |
gh pr edit ${{ github.event.pull_request.number }} \
--repo ${{ github.repository }} \
--remove-label "run-e2e-ci" 2>/dev/null || true

View File

@ -1,16 +1,12 @@
name: 'Integration test (Linux)'
on:
workflow_call:
inputs:
pr_sha:
description: 'PR head SHA (passed from parent workflow for workflow_dispatch)'
required: false
type: string
pr_number:
description: 'PR number (passed from parent workflow for workflow_dispatch)'
required: false
type: string
push:
branches: [ "main" ]
pull_request_target:
branches: [ "main" ]
types: [opened, synchronize, reopened, labeled]
jobs:
test-sglang-integration:
@ -21,23 +17,18 @@ jobs:
- name: trigger T-one test
if: ${{ env.tone_user_name != '' }}
run: |
# Priority: explicit inputs > PR event context > push SHA
SHA="${{ inputs.pr_sha || github.event.pull_request.head.sha || github.sha }}"
PR_ID="${{ inputs.pr_number || github.event.pull_request.number }}"
SHA="${{ github.event.pull_request.head.sha }}"
if [ "${{ github.event_name }}" = "push" ]; then
SHA="${{ github.sha }}"
PR_ID=""
fi
echo "PR_ID=${PR_ID}"
max_attempts=120
attempt=1
while [ $attempt -le $max_attempts ]; do
echo "Attempt $attempt: Fetching artifact..."
if curl -L -fs -o artifact.json -H "Accept: application/vnd.github+json" -H "X-GitHub-Api-Version: 2022-11-28" https://api.github.com/repos/${{ github.repository }}/actions/artifacts?per_page=100; then
if curl -L -fs -o artifact.json -H "Accept: application/vnd.github+json" -H "X-GitHub-Api-Version: 2022-11-28" https://api.github.com/repos/${{ github.repository }}/actions/artifacts; then
artifact_id=""
if jq empty artifact.json >/dev/null 2>&1; then
artifact_id=$(jq -r ".artifacts[] | select(.name | contains(\"py312\") ) | select(.name | contains(\"mooncake\") ) | select(.name | contains(\"cu130\") | not) | select(.workflow_run.head_sha == \"$SHA\" ) | .id" artifact.json | head -n 1)
artifact_id=$(jq -r ".artifacts[] | select(.name | contains(\"py312\") ) | select(.name | contains(\"cu130\") | not) | select(.workflow_run.head_sha == \"$SHA\" ) | .id" artifact.json | head -n 1)
else
echo "Failed to download artifact list. Retrying..."
fi
@ -62,14 +53,9 @@ jobs:
echo "Failed to fetch artifacts after $max_attempts attempts"
exit 1
fi
ENV_INFO="ARTIFACT_ID=${artifact_id} GIT_REPO=${{ github.repository }}"
if [ -n "$PR_ID" ]; then
ENV_INFO="${ENV_INFO} PR_ID=${PR_ID}"
fi
signature="${{ secrets.TONE_USER_NAME }}|${{ secrets.TONE_USER_TOKEN }}|$(python3 -c "import time;print(time.time())")"
signature="$(python3 -c "import base64;print(base64.b64encode(\"$signature\".encode('utf-8')).decode('utf-8'))")"
curl -s -H 'Content-Type: application/json' -X POST -d "{\"workspace\":\"mooncake_test\",\"project\":\"mooncake-ci\",\"template\":\"mooncake-ci-test\",\"name\":\"mooncake-ci-${SHA}\",\"username\":\"${{ secrets.TONE_USER_NAME }}\",\"env_ifs\":\" \",\"env_info\":\"${ENV_INFO}\",\"signature\":\"$signature\"}" https://tone.openanolis.cn/api/job/create/ > job.json
curl -s -H 'Content-Type: application/json' -X POST -d "{\"workspace\":\"mooncake_test\",\"project\":\"mooncake-ci\",\"template\":\"mooncake-ci-test\",\"name\":\"mooncake-ci-${SHA}\",\"username\":\"${{ secrets.TONE_USER_NAME }}\",\"env_ifs\":\" \",\"env_info\":\"ARTIFACT_ID=${artifact_id} GIT_REPO=${{ github.repository }}\",\"signature\":\"$signature\"}" https://tone.openanolis.cn/api/job/create/ > job.json
if [ "$(jq .code job.json)" == 200 ]; then
echo "job created"
else

View File

@ -18,6 +18,7 @@ jobs:
env:
BUILD_WITH_EP: "1"
CU13_BUILD: "1"
EP_TORCH_VERSIONS: "2.9.0;2.9.1;2.10.0"
TORCH_CUDA_ARCH_LIST: "8.0;9.0"
steps:
- name: Checkout source
@ -65,7 +66,7 @@ jobs:
sudo bash -x dependencies.sh -y
mkdir build
cd build
cmake .. -DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.9.0;2.9.1;2.10.0;2.11.0" -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
cmake .. -DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DWITH_EP=ON -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
shell: bash
- name: Build project

View File

@ -17,6 +17,7 @@ jobs:
python-version: ['3.10', '3.11', '3.12', '3.13']
env:
BUILD_WITH_EP: "1"
EP_TORCH_VERSIONS: "2.9.0;2.9.1;2.10.0"
TORCH_CUDA_ARCH_LIST: "8.0;9.0"
steps:
- name: Checkout source
@ -64,7 +65,7 @@ jobs:
sudo bash -x dependencies.sh -y
mkdir build
cd build
cmake .. -DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.9.0;2.9.1;2.10.0;2.11.0" -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
cmake .. -DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DWITH_EP=ON -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release
shell: bash
- name: Build project

1
.gitignore vendored
View File

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

4
.gitmodules vendored
View File

@ -2,7 +2,3 @@
path = extern/pybind11
url = https://github.com/pybind/pybind11.git
branch = stable
[submodule "extern/yalantinglibs"]
path = extern/yalantinglibs
url = https://github.com/alibaba/yalantinglibs.git
branch = v0.5.7

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

@ -14,14 +14,13 @@ endif()
option(WITH_TE "build mooncake transfer engine and sample code" ON)
option(WITH_STORE "build mooncake store library and sample code" ON)
option(WITH_STORE_GO "build Go bindings for mooncake store" OFF)
option(WITH_P2P_STORE "build p2p store library and sample code" OFF)
option(WITH_RUST_EXAMPLE "build the Rust interface and sample code for the transfer engine" OFF)
option(WITH_STORE_RUST "build the Rust bindings for the Mooncake Store" ON)
option(WITH_EP "build mooncake with expert parallelism support" OFF)
option(WITH_CONDUCTOR "build mooncake conductor and sample code" OFF)
include(${CMAKE_CURRENT_SOURCE_DIR}/mooncake-common/SetupPython.cmake)
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extern/pybind11)
set(PYTHON_EXECUTABLE "python3")
execute_process(
COMMAND ${PYTHON_EXECUTABLE} -c "import sys; print(sys.path[-1])"
OUTPUT_VARIABLE PYTHON_SYS_PATH
@ -41,25 +40,12 @@ option(STORE_USE_ETCD "build mooncake store with etcd" OFF)
if (STORE_USE_ETCD)
add_compile_definitions(STORE_USE_ETCD)
endif()
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)
@ -76,104 +62,19 @@ if (WITH_STORE)
include_directories(mooncake-store/include)
endif()
if (WITH_STORE_RUST)
if (NOT WITH_STORE)
message(FATAL_ERROR "WITH_STORE_RUST=ON requires WITH_STORE=ON")
endif()
message(STATUS "Mooncake Store Rust bindings will be built")
add_subdirectory(mooncake-store/rust)
endif()
option(EP_USE_IDE "Enable intelligent indexing for IDEs" OFF)
if (WITH_EP)
if (EP_USE_IDE)
message(WARNING "EP_USE_IDE enabled. DO NOT USE IN PRODUCTION!")
add_subdirectory(mooncake-ep)
include_directories(mooncake-ep/include)
add_subdirectory(mooncake-pg)
include_directories(mooncake-pg/include)
else ()
message(STATUS "WITH_EP enabled: building Mooncake EP and PG Python extensions")
find_package(CUDAToolkit REQUIRED)
message(STATUS "Detected CUDA version: ${CUDAToolkit_VERSION}")
# EP_TORCH_VERSIONS: semicolon-separated list of PyTorch versions to build for.
# Can be set via -DEP_TORCH_VERSIONS="2.9.1;2.8.0" or the EP_TORCH_VERSIONS env var.
# Empty means build with the currently-installed torch.
if(NOT EP_TORCH_VERSIONS)
set(EP_TORCH_VERSIONS "$ENV{EP_TORCH_VERSIONS}")
endif()
set(EP_TORCH_VERSIONS "${EP_TORCH_VERSIONS}" CACHE STRING
"PyTorch versions for EP/PG extensions, semicolon-separated (empty = use currently-installed torch)")
# TORCH_CUDA_ARCH_LIST forwarded to the torch CUDA extension build.
if(NOT TORCH_CUDA_ARCH_LIST)
set(TORCH_CUDA_ARCH_LIST "$ENV{TORCH_CUDA_ARCH_LIST}")
endif()
if(NOT TORCH_CUDA_ARCH_LIST)
set(TORCH_CUDA_ARCH_LIST "8.0;9.0")
endif()
set(TORCH_CUDA_ARCH_LIST "${TORCH_CUDA_ARCH_LIST}" CACHE STRING
"CUDA arch list for EP/PG extension builds (e.g. \"8.0;9.0\")")
# Staging directory: EP/PG .so files are placed here during make and later
# injected into the wheel AFTER auditwheel, so patchelf never touches the
# CUDA fatbins (which would cause cudaErrorInvalidKernelImage at runtime).
set(EP_PG_STAGING_DIR "${CMAKE_BINARY_DIR}/ep_pg_staging")
# Convert semicolon-separated lists to pipe-separated strings so they survive
# CMake's COMMAND list-splitting (semicolons are CMake list separators).
string(REPLACE ";" "|" _ep_torch_versions_pipe "${EP_TORCH_VERSIONS}")
string(REPLACE ";" "|" _torch_cuda_arch_list_pipe "${TORCH_CUDA_ARCH_LIST}")
add_custom_target(mooncake_ep_ext ALL
COMMAND ${CMAKE_COMMAND} -E make_directory "${EP_PG_STAGING_DIR}"
COMMAND ${CMAKE_COMMAND}
"-DSOURCE_DIR=${CMAKE_CURRENT_SOURCE_DIR}/mooncake-ep"
"-DEP_CUDA_MAJOR=${CUDAToolkit_VERSION_MAJOR}"
"-DEP_CUDA_MINOR=${CUDAToolkit_VERSION_MINOR}"
"-DEP_TORCH_VERSIONS=${_ep_torch_versions_pipe}"
"-DTORCH_CUDA_ARCH_LIST=${_torch_cuda_arch_list_pipe}"
"-DSTAGING_DIR=${EP_PG_STAGING_DIR}"
"-DENGINE_SO_PATH=$<TARGET_FILE:engine>"
-P "${CMAKE_CURRENT_SOURCE_DIR}/mooncake-ep/BuildEpExt.cmake"
COMMENT "Building Mooncake EP Python extension(s)"
DEPENDS engine
VERBATIM
)
add_custom_target(mooncake_pg_ext ALL
COMMAND ${CMAKE_COMMAND} -E make_directory "${EP_PG_STAGING_DIR}"
COMMAND ${CMAKE_COMMAND}
"-DSOURCE_DIR=${CMAKE_CURRENT_SOURCE_DIR}/mooncake-pg"
"-DEP_CUDA_MAJOR=${CUDAToolkit_VERSION_MAJOR}"
"-DEP_CUDA_MINOR=${CUDAToolkit_VERSION_MINOR}"
"-DEP_TORCH_VERSIONS=${_ep_torch_versions_pipe}"
"-DTORCH_CUDA_ARCH_LIST=${_torch_cuda_arch_list_pipe}"
"-DSTAGING_DIR=${EP_PG_STAGING_DIR}"
"-DENGINE_SO_PATH=$<TARGET_FILE:engine>"
-P "${CMAKE_CURRENT_SOURCE_DIR}/mooncake-pg/BuildPgExt.cmake"
COMMENT "Building Mooncake PG Python extension(s)"
DEPENDS engine mooncake_ep_ext
VERBATIM
)
endif ()
message(WARNING "Option `WITH_EP` is deprecated. Mooncake EP now builds with setuptools. Please set environment variable BUILD_WITH_EP=1 to enable.")
endif()
add_subdirectory(mooncake-integration)
if (WITH_STORE_GO AND WITH_STORE)
add_custom_target(build_store_go DEPENDS mooncake_store transfer_engine)
add_custom_command(
TARGET build_store_go
COMMAND bash build.sh ${CMAKE_BINARY_DIR} ${CMAKE_CURRENT_BINARY_DIR} ${USE_ETCD} ${USE_REDIS} ${USE_HTTP} ${USE_ETCD_LEGACY}
WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/mooncake-store/go
)
set_property(TARGET build_store_go PROPERTY EXCLUDE_FROM_ALL FALSE)
message(STATUS "Mooncake Store Go bindings will be built")
endif()
if (WITH_P2P_STORE)
add_subdirectory(mooncake-p2p-store)
message(STATUS "P2P Store will be built")
endif()
if (WITH_CONDUCTOR)
add_subdirectory(mooncake-conductor)
message(STATUS "Conductor will be built")
endif()

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

@ -15,8 +15,6 @@
[![Docs](https://img.shields.io/badge/docs-live-brightgreen)](https://kvcache-ai.github.io/Mooncake/)
[![PyPI](https://img.shields.io/pypi/v/mooncake-transfer-engine)](https://pypi.org/project/mooncake-transfer-engine)
[![PyPI - Python Version](https://img.shields.io/pypi/pyversions/mooncake-transfer-engine)](https://pypi.org/project/mooncake-transfer-engine)
[![CUDA <=12.9](https://img.shields.io/static/v1?label=CUDA&message=%3C%3D12.9&color=76B900)](https://pypi.org/project/mooncake-transfer-engine)
[![CUDA 13.0/13.1](https://img.shields.io/static/v1?label=CUDA&message=13.0%2F13.1&color=76B900)](https://pypi.org/project/mooncake-transfer-engine-cuda13)
[![PyPI - Downloads](https://img.shields.io/pypi/dm/mooncake-transfer-engine)](https://pypi.org/project/mooncake-transfer-engine)
[![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/kvcache-ai/Mooncake)
[![GitHub commit activity](https://img.shields.io/github/commit-activity/w/kvcache-ai/Mooncake)](https://github.com/kvcache-ai/Mooncake/graphs/commit-activity)
@ -31,10 +29,6 @@ This repository also hosts its technical report and the open-sourced traces.
<h2 id="updates">🔄 Updates</h2>
- **Mar 19, 2026**: [TorchSpec: Speculative Decoding Training at Scale](https://pytorch.org/blog/torchspec-speculative-decoding-training-at-scale) is [open sourced](https://github.com/torchspec-project/TorchSpec), using Mooncake to decouple inference and training via efficient hidden states management.
- **Mar 5, 2026**: [LightX2V](https://github.com/ModelTC/LightX2V/pull/893) now supports disaggregated deployment based on Mooncake, enabling encoder/transformer service decoupling with Mooncake Transfer Engine for high-performance cross-device and cross-machine data transfer.
- **Feb 25, 2026**: [SGLang](https://github.com/sgl-project/sglang) merged [Encoder Global Cache Manager](https://github.com/sgl-project/sglang/pull/16137), introducing a Mooncake-powered global multimodal embedding cache that enables cross-instance sharing of ViT embeddings to avoid redundant GPU computation.
- **Feb 24, 2026**: [vLLM-Omni](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/disaggregated_inference/) introduces disaggregated inference connectors with support for both `MooncakeStoreConnector` and `MooncakeTransferEngineConnector` for multi-node omni-modality pipelines.
- **Feb 12, 2026**: [Mooncake Joins PyTorch Ecosystem](https://pytorch.org/blog/mooncake-joins-pytorch-ecosystem/) We are thrilled to announce that Mooncake has officially joined the PyTorch Ecosystem!
- **Jan 28, 2026**: [FlexKV](https://github.com/taco-project/FlexKV), a distributed KV store and cache system from Tencent and NVIDIA in collaboration with the community, now supports [distributed KVCache reuse](https://github.com/taco-project/FlexKV/blob/main/docs/dist_reuse/README_en.md) with the Mooncake Transfer Engine.
- **Dec 27, 2025**: Collaboration with [ROLL](https://github.com/alibaba/ROLL)! Check out the paper [here](https://arxiv.org/abs/2512.22560).
@ -97,7 +91,7 @@ Mooncake establishes a full-stack, Tensor-oriented AI infrastructure where Tenso
### Use Transfer Engine Standalone ([Guide](https://kvcache-ai.github.io/Mooncake/design/transfer-engine/index.html))
Transfer Engine is a high-performance data transfer framework. Transfer Engine provides a unified interface to transfer data from DRAM, VRAM or NVMe, while the technical details related to hardware are hidden. Transfer Engine supports multiple communication protocols including TCP, RDMA (InfiniBand/RoCEv2/eRDMA/NVIDIA GPUDirect), NVMe over Fabric (NVMe-of), NVLink, HIP, CXL, and Ascend. When built with the corresponding runtime, Transfer Engine can also detect and route accelerator memory on CUDA, MUSA, HIP, and Cambricon MLU devices. For a complete list of supported protocols and configuration guide, see the [Supported Protocols Documentation](https://kvcache-ai.github.io/Mooncake/getting_started/supported-protocols.html).
Transfer Engine is a high-performance data transfer framework. Transfer Engine provides a unified interface to transfer data from DRAM, VRAM or NVMe, while the technical details related to hardware are hidden. Transfer Engine supports multiple communication protocols including TCP, RDMA (InfiniBand/RoCEv2/eRDMA/NVIDIA GPUDirect), NVMe over Fabric (NVMe-of), NVLink, HIP, CXL, and Ascend. For a complete list of supported protocols and configuration guide, see the [Supported Protocols Documentation](https://kvcache-ai.github.io/Mooncake/getting_started/supported-protocols.html).
#### Highlights
- **Efficient use of multiple RDMA NIC devices.** Transfer Engine supports the use of multiple RDMA NIC devices to achieve the *aggregation of transfer bandwidth*.
@ -178,7 +172,6 @@ The following need to be installed before running any component of Mooncake:
- RDMA Driver & SDK, such as Mellanox OFED.
- Python 3.10, virtual environment is recommended.
- CUDA 12.1 and above, including NVIDIA GPUDirect Storage Support, if the package is built with `-DUSE_CUDA` (disabled by default). *You may install them from [here](https://developer.nvidia.com/cuda-downloads)*.
- Cambricon Neuware, if the package is built with `-DUSE_MLU`. By default Mooncake looks for Neuware under `NEUWARE_HOME` or `/usr/local/neuware`.
### Use Python package
The simplest way to use Mooncake Transfer Engine is using `pip`:
@ -202,7 +195,6 @@ pip install mooncake-transfer-engine-non-cuda
> [!IMPORTANT]
> - The CUDA version (`mooncake-transfer-engine`) includes Mooncake-EP and GPU topology detection, requiring CUDA 12.1+.
> - The non-CUDA version (`mooncake-transfer-engine-non-cuda`) is for environments without CUDA dependencies.
> - MLU support is currently available through source builds with `-DUSE_MLU=ON`; there is no dedicated prebuilt MLU wheel yet.
> - If users encounter problems such as missing `lib*.so`, they should uninstall the package they installed and build the binaries manually.
### Use Docker image
@ -231,7 +223,6 @@ The following are additional dependencies for building Mooncake:
- Build essentials, including gcc, g++ (9.4+) and cmake (3.16+).
- Go 1.20+, if you want to build with `-DWITH_P2P_STORE`, `-DUSE_ETCD` (enabled by default to use etcd as metadata servers), or `-DSTORE_USE_ETCD` (use etcd for the failover of the store master).
- CUDA 12.1 and above, including NVIDIA GPUDirect Storage Support, if the package is built with `-DUSE_CUDA`. *This is NOT included in the `dependencies.sh` script. You may install them from [here](https://developer.nvidia.com/cuda-downloads)*.
- Cambricon Neuware, if you want to build with `-DUSE_MLU`. *This is NOT included in the `dependencies.sh` script.* Mooncake resolves it from `NEUWARE_HOME` or `/usr/local/neuware` by default, and also supports overriding `MLU_INCLUDE_DIR` / `MLU_LIB_DIR` during CMake configure.
- [Optional] Rust Toolchain, if you want to build with `-DWITH_RUST_EXAMPLE`. *This is NOT included in the `dependencies.sh` script.*
- [Optional] `hiredis`, if you want to build with `-DUSE_REDIS` to use Redis instead of etcd as metadata servers.
- [Optional] `curl`, if you want to build with `-DUSE_HTTP` to use HTTP instead of etcd as metadata servers.
@ -257,14 +248,6 @@ The build and installation steps are as follows:
sudo make install # optional, make it ready to be used by vLLM/SGLang
```
For Cambricon MLU builds, configure CMake with `-DUSE_MLU=ON`. For example:
```bash
mkdir build
cd build
cmake .. -DUSE_MLU=ON -DNEUWARE_ROOT=/usr/local/neuware
make -j
```
<h2 id="milestones"> 🛣️ Incoming Milestones</h2>

View File

@ -1,956 +0,0 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""
Mooncake KVCache Storage Benchmark Tool
"""
import argparse
import json
import time
import os
import statistics
import random
import errno
from pathlib import Path
from typing import Dict, List, Optional
from dataclasses import dataclass
# ============================================================================
# Constants
# ============================================================================
BLOCK_SIZE_TOKENS = 512 # Number of tokens per block
DEFAULT_BYTES_PER_TOKEN = 2048 # 7B model FP16 (2KB per token)
BLOCK_SIZE_BYTES = BLOCK_SIZE_TOKENS * DEFAULT_BYTES_PER_TOKEN # 1MB per block
MIN_LATENCY_MS = 0.001 # Minimum latency in milliseconds (1 microsecond)
# Model KVCache sizes (bytes per token, based on LMCache calculator)
# Source: https://lmcache.ai/kv_cache_calculator.html
MODEL_BYTES_PER_TOKEN = {
"llama-3.1-405b": 327680,
"qwen3-32b": 81920,
"deepseek-v3": 1748992,
"glm-4.6": 157013,
"default": DEFAULT_BYTES_PER_TOKEN,
}
# ============================================================================
# Data Structures
# ============================================================================
@dataclass
class KVCacheRequest:
"""KVCache request
Attributes:
timestamp: Request timestamp in milliseconds
hash_ids: List of block IDs (each ID corresponds to a 512-token block)
input_length: Input token count
output_length: Output token count
"""
timestamp: float
hash_ids: List[int]
input_length: int
output_length: int
# ============================================================================
# Storage Layer: Offset Allocator
# ============================================================================
class OffsetAllocatorStorage:
"""High-performance block storage based on Offset Allocator
Architecture:
-----------
1. Single large file stores all blocks (avoids file explosion)
2. Uses offset to manage file space (similar to Mooncake's OffsetAllocator)
3. hash_id -> offset mapping stored in memory (fast lookup)
Block Organization:
-----------
Each block corresponds to 512 tokens, fixed size 1MB:
- hash_id[0] -> block_0 (tokens [0...511]) -> offset 0
- hash_id[1] -> block_1 (tokens [512...1023]) -> offset 1
- hash_id[i] -> block_i (tokens [i*512...(i+1)*512-1]) -> offset i
Performance Advantages:
-----------
- Only one file, no file explosion
- Offset reuse, reduces memory allocation
- pread/pwrite, thread-safe, no seek needed
- Keep fd open, reduces open/close overhead
- Metadata in memory, O(1) lookup
Attributes:
storage_dir: Storage directory path
block_size_bytes: Block size in bytes
max_blocks: Maximum number of blocks
hash_id_to_offset: hash_id -> offset mapping
free_offsets: List of reusable offsets
next_offset: Next allocatable offset
"""
def __init__(self, storage_dir: str, bytes_per_token: int = DEFAULT_BYTES_PER_TOKEN,
max_blocks: int = 100000, block_size_tokens: int = 512,
fsync_mode: str = 'batch', fsync_batch_size: int = 100):
"""Initialize Offset Allocator storage
Args:
storage_dir: Storage directory path
bytes_per_token: Bytes per token
max_blocks: Maximum number of blocks (determines file size)
block_size_tokens: Number of tokens per block
fsync_mode: When to fsync ('batch', 'always', 'end', 'none')
fsync_batch_size: Number of writes between fsync in batch mode
"""
self.storage_dir = Path(storage_dir)
self.bytes_per_token = bytes_per_token
self.block_size_tokens = block_size_tokens
self.block_size_bytes = self.block_size_tokens * self.bytes_per_token
self.max_blocks = max_blocks
# Fsync configuration
self.fsync_mode = fsync_mode
self.fsync_batch_size = fsync_batch_size
self.pending_sync_count = 0
# Create storage directory
self.storage_dir.mkdir(parents=True, exist_ok=True)
# Single large file
self.storage_file = self.storage_dir / "kvcache_storage.bin"
self.file_size = self.max_blocks * self.block_size_bytes
# Initialize storage file
if not self.storage_file.exists():
self._init_storage_file()
# hash_id -> offset mapping (metadata, in memory)
self.hash_id_to_offset: Dict[int, int] = {}
# Offset allocator (free list)
self.free_offsets: List[int] = []
self.next_offset = 0
# File descriptor (keep open, avoid repeated open/close)
self.fd = None
# Pre-allocated data buffer with pattern to avoid SSD compression artifacts
# Using a repeating pattern that looks like realistic data (not all zeros)
# Pattern: 64-byte repeated sequence mixed with some variation
pattern = bytes([(i & 0xFF) for i in range(256)]) # 0-255 byte pattern
pattern_repeats = (self.block_size_bytes // len(pattern)) + 1
self._data_buffer = (pattern * pattern_repeats)[:self.block_size_bytes]
# Statistics
self.stats = {
'read_count': 0,
'write_count': 0,
'read_bytes': 0,
'write_bytes': 0,
'read_latencies_ms': [],
'write_latencies_ms': [],
'sync_count': 0, # Number of fsync operations performed
}
# ========================================================================
# Internal Methods
# ========================================================================
def _init_storage_file(self):
"""Initialize storage file (pre-allocate space)
Create sparse file to avoid actual disk space usage until data is written
"""
with open(self.storage_file, 'wb') as f:
f.seek(self.file_size - 1)
f.write(b'\0')
f.flush()
os.fsync(f.fileno())
def _get_fd(self):
"""Get file descriptor (lazy open)
Returns:
int: File descriptor
"""
if self.fd is None:
# Use O_RDWR | O_CREAT, no O_DIRECT (Python compatibility)
self.fd = os.open(self.storage_file, os.O_RDWR | os.O_CREAT)
return self.fd
def _allocate_offset(self) -> int:
"""Allocate a new offset
Prioritize reusing freed offsets, otherwise allocate new offset
Returns:
int: Allocated offset
"""
if self.free_offsets:
return self.free_offsets.pop()
offset = self.next_offset
self.next_offset += 1
return offset
def _free_offset(self, offset: int):
"""Free offset for reuse
Args:
offset: Offset to free
"""
self.free_offsets.append(offset)
# ========================================================================
# Public Interface
# ========================================================================
def block_exists(self, hash_id: int) -> bool:
"""Check if block exists
Args:
hash_id: Unique block identifier
Returns:
bool: Whether block exists
"""
return hash_id in self.hash_id_to_offset
def read_block(self, hash_id: int) -> float:
"""Read block using pread
Args:
hash_id: Unique block identifier
Returns:
float: Read latency in milliseconds, or 0 if block doesn't exist
"""
if hash_id not in self.hash_id_to_offset:
return 0.0 # Block doesn't exist, no latency to measure
offset = self.hash_id_to_offset[hash_id]
file_offset = offset * self.block_size_bytes
start = time.perf_counter()
try:
fd = self._get_fd()
data = os.pread(fd, self.block_size_bytes, file_offset)
latency_ms = (time.perf_counter() - start) * 1000.0
self.stats['read_count'] += 1
self.stats['read_bytes'] += len(data)
self.stats['read_latencies_ms'].append(latency_ms)
return latency_ms
except OSError as e:
print(f"Error reading block {hash_id} at offset {file_offset}: {e}")
return 0.0 # Error case, don't pollute stats
def write_block(self, hash_id: int) -> float:
"""Write block using pwrite
Args:
hash_id: Unique block identifier
Returns:
float: Write latency in milliseconds
"""
# Allocate offset
offset = self._allocate_offset()
file_offset = offset * self.block_size_bytes
# Use pre-allocated buffer (much faster than os.urandom)
data = self._data_buffer
start = time.perf_counter()
try:
fd = self._get_fd()
written = os.pwrite(fd, data, file_offset)
write_done = time.perf_counter()
# Conditional fsync based on mode
if self.fsync_mode == 'always':
# Include fsync in latency measurement
os.fsync(fd)
self.stats['sync_count'] += 1
self.pending_sync_count = 0
latency_ms = (time.perf_counter() - start) * 1000.0
# Evict from page cache AFTER fsync to ensure reads measure actual SSD performance
os.posix_fadvise(fd, file_offset, self.block_size_bytes, os.POSIX_FADV_DONTNEED)
elif self.fsync_mode == 'batch':
# For batch mode, only measure write time (fsync is deferred)
self.pending_sync_count += 1
if self.pending_sync_count >= self.fsync_batch_size:
os.fsync(fd)
self.stats['sync_count'] += 1
self.pending_sync_count = 0
latency_ms = (write_done - start) * 1000.0 # Only write time
# Evict from page cache after each write
os.posix_fadvise(fd, file_offset, self.block_size_bytes, os.POSIX_FADV_DONTNEED)
elif self.fsync_mode == 'none':
latency_ms = (write_done - start) * 1000.0
# Evict from page cache even when not syncing
os.posix_fadvise(fd, file_offset, self.block_size_bytes, os.POSIX_FADV_DONTNEED)
else: # 'end' mode
latency_ms = (write_done - start) * 1000.0
# Evict from page cache (fsync will happen at the end)
os.posix_fadvise(fd, file_offset, self.block_size_bytes, os.POSIX_FADV_DONTNEED)
# Update mapping
self.hash_id_to_offset[hash_id] = offset
self.stats['write_count'] += 1
self.stats['write_bytes'] += written
self.stats['write_latencies_ms'].append(latency_ms)
return latency_ms
except OSError as e:
if e.errno == errno.ENOSPC:
print(f"Error: Disk full when writing block {hash_id} at offset {file_offset}")
else:
print(f"Error writing block {hash_id} at offset {file_offset}: {e}")
return 0.0 # Error case, don't pollute stats
def __enter__(self):
"""Context manager entry"""
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""Context manager exit - ensures cleanup"""
# Perform final fsync before closing for 'end' and 'batch' modes
self._finalize_sync()
self.close(force_sync=False) # Already synced above
return False
def _finalize_sync(self):
"""Perform final fsync before closing (for 'end' mode and pending batch writes)"""
if self.fd is not None:
if self.fsync_mode == 'end':
try:
os.fsync(self.fd)
self.stats['sync_count'] += 1
except OSError:
pass
elif self.fsync_mode == 'batch' and self.pending_sync_count > 0:
# Flush remaining pending writes
try:
os.fsync(self.fd)
self.stats['sync_count'] += 1
self.pending_sync_count = 0
except OSError:
pass
def close(self, force_sync: bool = True):
"""Close file
Args:
force_sync: Whether to force fsync before closing
"""
# For backward compatibility with non-context-manager usage
if force_sync:
self._finalize_sync()
if self.fd is not None:
os.close(self.fd)
self.fd = None
def get_stats(self) -> Dict:
"""Get statistics
Returns:
Dict: Dictionary containing read/write statistics
"""
def calc_stats(latencies):
"""Calculate latency statistics"""
if not latencies:
return {'avg_ms': 0, 'p50_ms': 0, 'p95_ms': 0, 'p99_ms': 0}
return {
'avg_ms': statistics.mean(latencies),
**calc_percentiles(latencies),
}
return {
'read': {
'count': self.stats['read_count'],
'mb': self.stats['read_bytes'] / 1024 / 1024,
**calc_stats(self.stats['read_latencies_ms'])
},
'write': {
'count': self.stats['write_count'],
'mb': self.stats['write_bytes'] / 1024 / 1024,
**calc_stats(self.stats['write_latencies_ms'])
},
'sync_count': self.stats['sync_count'],
'total_blocks': len(self.hash_id_to_offset),
'free_blocks': len(self.free_offsets),
}
# ============================================================================
# Benchmark Layer
# ============================================================================
class StorageBenchmark:
"""KVCache storage benchmark
Based on Mooncake OffsetAllocator + vLLM PagedAttention implementation:
Example:
-----
Request A: [1, 2, 4]
-> hash_id 1 -> not exist, write block_1 (offset=0, 1MB)
-> hash_id 2 -> not exist, write block_2 (offset=1, 1MB)
-> hash_id 4 -> not exist, write block_4 (offset=2, 1MB)
Request B: [1, 2, 4, 6]
-> hash_id 1 -> exists, read block_1 (offset=0) prefix reuse
-> hash_id 2 -> exists, read block_2 (offset=1) prefix reuse
-> hash_id 4 -> exists, read block_4 (offset=2) prefix reuse
-> hash_id 6 -> not exist, write block_6 (offset=3, 1MB)
Performance Advantages:
---------
- Single file operation, no file explosion
- Offset reuse, reduces memory allocation
- pread/pwrite, thread-safe
"""
def __init__(self, storage_dir: str, bytes_per_token: int = DEFAULT_BYTES_PER_TOKEN,
max_blocks: int = 100000, block_size_tokens: int = 512,
fsync_mode: str = 'batch', fsync_batch_size: int = 100):
"""Initialize benchmark
Args:
storage_dir: Storage directory
bytes_per_token: Bytes per token
max_blocks: Maximum number of blocks
block_size_tokens: Number of tokens per block
fsync_mode: When to fsync ('batch', 'always', 'end', 'none')
fsync_batch_size: Number of writes between fsync in batch mode
"""
self.storage = OffsetAllocatorStorage(
storage_dir, bytes_per_token, max_blocks,
block_size_tokens, fsync_mode, fsync_batch_size
)
self.bytes_per_token = bytes_per_token
self.block_size_tokens = block_size_tokens
# Statistics
self.stats = {
'total_requests': 0,
'total_blocks': 0,
'read_blocks': 0,
'write_blocks': 0,
'prefix_hit_blocks': 0, # Number of prefix hit blocks
'request_latencies_ms': [],
}
def process_request(self, req: KVCacheRequest) -> float:
"""Process a KVCache request
Based on vLLM's prefix caching mechanism:
- Each hash_id corresponds to an independent block
- Prefix reuse achieved through hash_id matching
Args:
req: KVCache request
Returns:
float: Request latency in milliseconds
"""
self.stats['total_requests'] += 1
self.stats['total_blocks'] += len(req.hash_ids)
start_time = time.perf_counter()
total_latency = 0.0
# Process each hash_id (in order)
for hash_id in req.hash_ids:
if self.storage.block_exists(hash_id):
# Block exists, read (reuse cached block)
total_latency += self.storage.read_block(hash_id)
self.stats['read_blocks'] += 1
self.stats['prefix_hit_blocks'] += 1 # Count all cache hits as prefix reuse
else:
# Block doesn't exist, write (new block)
total_latency += self.storage.write_block(hash_id)
self.stats['write_blocks'] += 1
latency_ms = total_latency if total_latency > 0 else MIN_LATENCY_MS
self.stats['request_latencies_ms'].append(latency_ms)
return latency_ms
def get_stats(self) -> Dict:
"""Get statistics
Returns:
Dict: Statistics dictionary
"""
storage_stats = self.storage.get_stats()
request_latencies = self.stats['request_latencies_ms']
if request_latencies:
latency_stats = {
'avg_ms': statistics.mean(request_latencies),
**calc_percentiles(request_latencies),
}
else:
latency_stats = {'avg_ms': 0, 'p50_ms': 0, 'p95_ms': 0, 'p99_ms': 0}
total_blocks = self.stats['total_blocks']
read_blocks = self.stats['read_blocks']
write_blocks = self.stats['write_blocks']
return {
'total_requests': self.stats['total_requests'],
'total_blocks': total_blocks,
'read_blocks': read_blocks,
'write_blocks': write_blocks,
'prefix_hit_blocks': self.stats['prefix_hit_blocks'],
'block_hit_rate': read_blocks / total_blocks if total_blocks > 0 else 0,
'write_ratio': write_blocks / total_blocks if total_blocks > 0 else 0,
'tokens_per_block': self.block_size_tokens, # Configurable block size in tokens
'latency': latency_stats,
'storage': storage_stats,
}
def __enter__(self):
"""Context manager entry"""
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""Context manager exit - ensures cleanup"""
self.close()
return False
def close(self, force_sync: bool = True):
"""Close storage
Args:
force_sync: Whether to force final sync before closing
"""
self.storage.close(force_sync=force_sync)
# ============================================================================
# Utility Functions
# ============================================================================
def calc_percentiles(data: List[float]) -> Dict[str, float]:
"""Calculate latency percentiles
Uses linear interpolation for accurate percentile calculation.
This is more accurate than statistics.quantiles() for small datasets.
Args:
data: List of latency values in milliseconds
Returns:
Dict containing p50, p95, p99 percentiles
"""
if not data:
return {'p50_ms': 0, 'p95_ms': 0, 'p99_ms': 0}
# Sort data for percentile calculation
sorted_data = sorted(data)
n = len(sorted_data)
def get_percentile(p: float) -> float:
"""Get percentile using linear interpolation
Args:
p: Percentile (0-100)
Returns:
Value at percentile
"""
index = (n - 1) * p / 100
lower = int(index)
upper = min(lower + 1, n - 1)
if lower == upper:
return sorted_data[lower]
# Linear interpolation
weight = index - lower
return sorted_data[lower] * (1 - weight) + sorted_data[upper] * weight
return {
'p50_ms': get_percentile(50),
'p95_ms': get_percentile(95),
'p99_ms': get_percentile(99),
}
# ============================================================================
# Trace Loader
# ============================================================================
class TraceLoader:
"""Load KVCache trace"""
def __init__(self, trace_path: str):
"""Initialize trace loader
Args:
trace_path: Trace file path
"""
self.trace_path = trace_path
self.requests = []
self._load_trace()
def _load_trace(self):
"""Load trace file with error handling"""
line_num = 0
try:
with open(self.trace_path, 'r') as f:
for line in f:
line_num += 1
line = line.strip()
if not line:
continue
try:
req = json.loads(line)
# Validate required fields
if not all(k in req for k in ['timestamp', 'hash_ids', 'input_length', 'output_length']):
print(f"Warning: Line {line_num} missing required fields, skipping")
continue
if not isinstance(req['hash_ids'], list):
print(f"Warning: Line {line_num} has invalid hash_ids (not a list), skipping")
continue
self.requests.append(KVCacheRequest(
timestamp=float(req['timestamp']),
hash_ids=req['hash_ids'],
input_length=int(req['input_length']),
output_length=int(req['output_length'])
))
except (json.JSONDecodeError, ValueError, KeyError) as e:
print(f"Warning: Line {line_num} has invalid format: {e}, skipping")
continue
except FileNotFoundError:
raise FileNotFoundError(f"Trace file not found: {self.trace_path}")
except OSError as e:
raise OSError(f"Error reading trace file {self.trace_path}: {e}")
def get_requests(self) -> List[KVCacheRequest]:
"""Get request list
Returns:
List[KVCacheRequest]: Request list
"""
return self.requests
# ============================================================================
# Benchmark Runner
# ============================================================================
def run_benchmark(trace_path: str, storage_dir: str, bytes_per_token: int = DEFAULT_BYTES_PER_TOKEN,
max_requests: Optional[int] = None, max_blocks: int = 100000,
replay_timestamps: bool = False, time_scale: float = 1.0,
block_size_tokens: int = 512,
fsync_mode: str = 'batch', fsync_batch_size: int = 100) -> Dict:
"""Run benchmark
Args:
trace_path: Trace file path
storage_dir: Storage directory
bytes_per_token: Bytes per token
max_requests: Maximum number of requests (None = all)
max_blocks: Maximum number of blocks
replay_timestamps: Whether to replay timestamps from trace (simulate realistic timing)
time_scale: Time scaling factor (1.0=real-time, 0.1=10x speed, 10.0=0.1x speed)
block_size_tokens: Number of tokens per block
fsync_mode: When to fsync ('batch', 'always', 'end', 'none')
fsync_batch_size: Number of writes between fsync in batch mode
Returns:
Dict: Benchmark results
"""
block_size_bytes = block_size_tokens * bytes_per_token
print(f"\n{'='*80}")
print(f"Running: {Path(trace_path).name}")
print(f"Architecture: Offset Allocator (Mooncake style)")
print(f"Block size: {block_size_tokens} tokens/block ({block_size_bytes:,} bytes)")
print(f"Storage: Single large file with offset-based block management")
print(f"Bytes per token: {bytes_per_token}")
print(f"Max blocks: {max_blocks}")
print(f"Fsync mode: {fsync_mode}" + (f" (batch_size={fsync_batch_size})" if fsync_mode == 'batch' else ''))
print(f"Timestamp replay: {'Enabled' if replay_timestamps else 'Disabled'}")
if replay_timestamps:
scale_desc = 'real-time' if time_scale == 1.0 else f'{1/time_scale:.1f}x speed' if time_scale < 1.0 else f'{time_scale}x slower'
print(f"Time scale: {time_scale}x ({scale_desc})")
print(f"{'='*80}")
# Load trace
loader = TraceLoader(trace_path)
requests = loader.get_requests()
if max_requests:
requests = requests[:max_requests]
print(f"Loaded {len(requests)} requests")
# Show timestamp range
if replay_timestamps and requests:
timestamps = [req.timestamp for req in requests]
time_span_ms = max(timestamps) - min(timestamps)
print(f"Timestamp range: {min(timestamps):.1f} - {max(timestamps):.1f} ms (span: {time_span_ms:.1f} ms)")
# Create benchmark instance with context manager for cleanup
with StorageBenchmark(
storage_dir, bytes_per_token, max_blocks,
block_size_tokens, fsync_mode, fsync_batch_size
) as benchmark:
# Run benchmark
start_time = time.perf_counter()
total_io_time = 0.0 # Actual I/O time (excluding sleep)
last_timestamp = None
base_time = time.time() # Use wall time for replay synchronization
for i, req in enumerate(requests):
# Replay by timestamps
sleep_time = 0.0
if replay_timestamps and last_timestamp is not None:
# Calculate time interval from previous request
delta_ms = req.timestamp - last_timestamp
sleep_time = delta_ms / 1000.0 / time_scale # Apply time scaling
if sleep_time > 0:
time.sleep(sleep_time)
# Process request (measure I/O time)
req_start = time.perf_counter()
benchmark.process_request(req)
req_io_time = time.perf_counter() - req_start
total_io_time += req_io_time
# Record current request timestamp
last_timestamp = req.timestamp
# Progress output
if (i + 1) % 100 == 0:
if replay_timestamps:
elapsed_wall_time = time.time() - base_time
simulated_time = (req.timestamp - requests[0].timestamp) / 1000.0 / time_scale
print(f" Processed {i + 1}/{len(requests)}... (wall: {elapsed_wall_time:.1f}s, simulated: {simulated_time:.1f}s, io: {total_io_time:.1f}s)")
else:
print(f" Processed {i + 1}/{len(requests)}...")
elapsed = time.perf_counter() - start_time
# Perform final sync to include it in stats
benchmark.storage._finalize_sync()
# Get statistics (context manager will handle cleanup)
stats = benchmark.get_stats()
# Calculate actual I/O time (excluding sleep)
io_time = total_io_time if replay_timestamps else elapsed
return {
'trace_file': Path(trace_path).name,
'total_requests': len(requests),
'simulation_time_s': elapsed,
'io_time_s': io_time, # Actual I/O time
'wall_time_s': elapsed, # Wall time (including sleep)
'requests_per_second': len(requests) / io_time if io_time > 0 else 0, # Based on I/O time
'timestamp_replay_enabled': replay_timestamps,
'time_scale': time_scale,
'bytes_per_token': bytes_per_token,
'block_size_tokens': block_size_tokens,
'fsync_mode': fsync_mode,
**stats,
}
# ============================================================================
# Result Output
# ============================================================================
def print_results(results: List[Dict]):
"""Print benchmark results
Args:
results: List of benchmark results
"""
for i, r in enumerate(results, 1):
print(f"\n{'='*80}")
print(f" [{i}/{len(results)}] {r['trace_file']}")
print(f"{'='*80}")
print(f"\n[Performance Overview]")
print(f" Total Requests: {r['total_requests']:,}")
print(f" Queries Per Second (QPS): {r['requests_per_second']:.2f}")
print(f" Cache Hit Rate: {r['block_hit_rate']:.2%}")
print(f" Write Ratio: {r['write_ratio']:.2%}")
print(f" Total Blocks: {r['total_blocks']:,}")
print(f" Read Blocks: {r['read_blocks']:,}")
print(f" Write Blocks: {r['write_blocks']:,}")
print(f" Prefix Hits: {r['prefix_hit_blocks']:,}")
print(f"\n[Latency Analysis]")
req_lat = r['latency']
print(f" Request Latency (End-to-End): Avg={req_lat['avg_ms']:.2f}ms, P50={req_lat['p50_ms']:.2f}ms, P95={req_lat['p95_ms']:.2f}ms, P99={req_lat['p99_ms']:.2f}ms")
read_lat = r['storage']['read']
write_lat = r['storage']['write']
print(f" Single I/O Operation (Per Block):")
print(f" Read: Avg={read_lat.get('avg_ms', 0):.3f}ms, P50={read_lat.get('p50_ms', 0):.3f}ms, P95={read_lat.get('p95_ms', 0):.3f}ms, P99={read_lat.get('p99_ms', 0):.3f}ms")
print(f" Write: Avg={write_lat.get('avg_ms', 0):.3f}ms, P50={write_lat.get('p50_ms', 0):.3f}ms, P95={write_lat.get('p95_ms', 0):.3f}ms, P99={write_lat.get('p99_ms', 0):.3f}ms")
print(f"\n[I/O & Bandwidth]")
print(f" Total Read I/O: {r['storage']['read']['mb']:>10.1f} MB ({r['storage']['read']['count']:,} ops)")
print(f" Total Write I/O: {r['storage']['write']['mb']:>10.1f} MB ({r['storage']['write']['count']:,} ops)")
io_time = r['io_time_s']
bandwidth = (r['storage']['read']['mb'] + r['storage']['write']['mb']) / io_time
print(f" Effective Bandwidth: {bandwidth:>10.1f} MB/s")
print(f"\n[Storage Details]")
print(f" Blocks in Use: {r['storage']['total_blocks']:>10,}")
print(f" Free Blocks: {r['storage']['free_blocks']:>10,}")
print(f" Tokens per Block: {r['tokens_per_block']:>10,}")
print(f" Block Size: {r['tokens_per_block'] * r.get('bytes_per_token', 2048) / 1024 / 1024:>10.2f} MB")
if 'sync_count' in r['storage']:
print(f" Fsync Operations: {r['storage']['sync_count']:>10,}")
print(f"\n[Execution Time]")
if r.get('timestamp_replay_enabled'):
print(f" Wall Time (Total): {r['wall_time_s']:>10.2f} s")
print(f" I/O Time (Actual): {r['io_time_s']:>10.2f} s")
print(f" Sleep Time (Replay): {r['wall_time_s'] - r['io_time_s']:>10.2f} s")
else:
print(f" Total Execution Time: {r['wall_time_s']:>10.2f} s")
print(f"\n{'='*80}\n")
# ============================================================================
# Main Program
# ============================================================================
def main():
"""Main entry point"""
parser = argparse.ArgumentParser(
description='Mooncake KVCache Storage Benchmark',
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
# Quick test (100 requests)
python storage_benchmark.py --scenario=toolagent --max-requests=100
# Test with large model preset (Llama-3.1-405B)
python storage_benchmark.py --scenario=toolagent --model=llama-3.1-405b --max-requests=100
# Test with Deepseek V3 (extra large model)
python storage_benchmark.py --scenario=toolagent --model=deepseek-v3 --max-requests=100
# Realistic replay (with timestamps, 10x speed)
python storage_benchmark.py --scenario=toolagent --max-requests=1000 \\
--replay-timestamps --time-scale=0.1
# All scenarios with custom bytes_per_token
python storage_benchmark.py --scenario=all --bytes-per-token=512
# Test with different block sizes and fsync modes
python storage_benchmark.py --scenario=toolagent --block-size-tokens=256 --fsync-mode=always
# Test with custom fsync batch size
python storage_benchmark.py --scenario=toolagent --fsync-mode=batch --fsync-batch-size=50
Performance Tuning:
--fsync-mode=batch (default): Balance between performance and safety
--fsync-mode=always: Safest but slowest, measures full persistence cost
--fsync-mode=end: Fastest, only measures write I/O (not persistence)
--fsync-mode=none: Testing only, no durability guarantees
Available model presets:
llama-3.1-405b, qwen3-32b, deepseek-v3, glm-4.6, default
For more information: tools/STORAGE_BENCHMARK_README.md
"""
)
parser.add_argument('--trace-dir', type=str, default='../../FAST25-release/traces',
help='Trace files directory')
parser.add_argument('--scenario', type=str, choices=['conversation', 'synthetic', 'toolagent', 'all'],
default='toolagent', help='Test scenario')
parser.add_argument('--storage-dir', type=str, default='/tmp/mooncake_bench',
help='Storage directory')
parser.add_argument('--model', type=str, choices=list(MODEL_BYTES_PER_TOKEN.keys()),
default='default',
help=f'Model preset (overrides --bytes-per-token). Available: {", ".join(MODEL_BYTES_PER_TOKEN.keys())}')
parser.add_argument('--bytes-per-token', type=int, default=DEFAULT_BYTES_PER_TOKEN,
help='Bytes per token (default %d, overridden by --model if specified)' % DEFAULT_BYTES_PER_TOKEN)
parser.add_argument('--max-requests', type=int, default=None,
help='Maximum number of requests (default: unlimited)')
parser.add_argument('--max-blocks', type=int, default=100000,
help='Maximum number of blocks in storage file (determines file size)')
parser.add_argument('--replay-timestamps', action='store_true',
help='Enable timestamp replay (simulate realistic request timing)')
parser.add_argument('--time-scale', type=float, default=1.0,
help='Time scaling factor (1.0=real-time, 0.1=10x speed, 10.0=0.1x speed)')
parser.add_argument('--block-size-tokens', type=int, default=512,
help='Number of tokens per block (default: 512)')
parser.add_argument('--fsync-mode', type=str, choices=['batch', 'always', 'end', 'none'],
default='batch',
help='When to fsync: batch=every N writes (default), always=after each write, end=only at close, none=never')
parser.add_argument('--fsync-batch-size', type=int, default=100,
help='Number of writes between fsync in batch mode (default: 100)')
args = parser.parse_args()
# Print benchmark header
print(f"\n{'='*80}")
print(f"{'Mooncake KVCache Storage Benchmark':^80}")
print(f"{'='*80}")
# Determine bytes_per_token (model preset takes precedence)
bytes_per_token = MODEL_BYTES_PER_TOKEN.get(args.model, args.bytes_per_token)
if args.model != 'default':
print(f"Using model preset: {args.model} ({bytes_per_token} bytes/token, ~{bytes_per_token/1024:.1f} KB/token)")
else:
print(f"Using custom bytes_per_token: {bytes_per_token}")
# Determine test scenarios
scenarios = ['conversation', 'synthetic', 'toolagent'] if args.scenario == 'all' else [args.scenario]
trace_files = {
'conversation': 'conversation_trace.jsonl',
'synthetic': 'synthetic_trace.jsonl',
'toolagent': 'toolagent_trace.jsonl'
}
# Run benchmarks
results = []
for scenario in scenarios:
trace_path = Path(args.trace_dir) / trace_files[scenario]
if trace_path.exists():
result = run_benchmark(
str(trace_path),
str(Path(args.storage_dir) / scenario),
bytes_per_token,
args.max_requests,
args.max_blocks,
args.replay_timestamps,
args.time_scale,
args.block_size_tokens,
args.fsync_mode,
args.fsync_batch_size
)
results.append(result)
else:
print(f"Warning: Trace file not found: {trace_path}")
# Print results
if results:
print_results(results)
if __name__ == '__main__':
main()

View File

@ -23,7 +23,8 @@ NC="\033[0m" # No Color
# Configuration
REPO_ROOT=`pwd`
GITHUB_PROXY=${GITHUB_PROXY:-"https://github.com"}
GOVER=1.25.9
GOVER=1.23.8
YALANTINGLIBS_VERSION=0.5.7
# Function to print section headers
print_section() {
@ -75,7 +76,8 @@ echo -e "${YELLOW}Mooncake Dependencies Installer${NC}"
echo -e "This script will install all required dependencies for Mooncake."
echo -e "The following components will be installed:"
echo -e " - System packages (build tools, libraries)"
echo -e " - Git submodules (including pybind11 and yalantinglibs)"
echo -e " - yalantinglibs"
echo -e " - Git submodules"
echo -e " - Go $GOVER"
echo
@ -101,7 +103,6 @@ echo -e "${YELLOW}This may take a few minutes...${NC}"
SYSTEM_PACKAGES="build-essential \
cmake \
ninja-build \
git \
wget \
unzip \
@ -126,7 +127,7 @@ SYSTEM_PACKAGES="build-essential \
libmsgpack-dev \
libzstd-dev \
libasio-dev \
libxxhash-dev \
libzmq3-dev \
pkg-config \
patchelf \
libc6-dev \
@ -136,34 +137,48 @@ apt-get install -y $SYSTEM_PACKAGES
check_success "Failed to install system packages"
print_success "System packages installed successfully"
# Initialize and update git submodules
print_section "Initializing Git Submodules"
# Install yalantinglibs
print_section "Installing yalantinglibs"
# Check if .gitmodules exists
if [ -f "${REPO_ROOT}/.gitmodules" ]; then
echo "Enter repository root: ${REPO_ROOT}"
cd "${REPO_ROOT}"
check_success "Failed to change to repository root directory"
echo "Initializing git submodules..."
git submodule sync --recursive
check_success "Failed to sync git submodules"
git submodule update --init --recursive
check_success "Failed to initialize git submodules"
print_success "Git submodules initialized and updated successfully"
else
echo -e "${YELLOW}No .gitmodules file found. Skipping...${NC}"
exit 1
# Check if thirdparties directory exists
if [ ! -d "${REPO_ROOT}/thirdparties" ]; then
mkdir -p "${REPO_ROOT}/thirdparties"
check_success "Failed to create thirdparties directory"
fi
# Build and install yalantinglibs from submodule
print_section "Installing yalantinglibs"
cd "${REPO_ROOT}/extern/yalantinglibs"
check_success "Failed to change to yalantinglibs submodule directory"
# Change to thirdparties directory
cd "${REPO_ROOT}/thirdparties"
check_success "Failed to change to thirdparties directory"
# Check if yalantinglibs is already installed
if [ -d "yalantinglibs-${YALANTINGLIBS_VERSION}" ]; then
echo -e "${YELLOW}yalantinglibs-${YALANTINGLIBS_VERSION} directory already exists. Removing for fresh install...${NC}"
rm -rf yalantinglibs-${YALANTINGLIBS_VERSION}
check_success "Failed to remove existing yalantinglibs directory"
fi
# Download yalantinglibs
YALANTINGLIBS_ZIPFILE="yalantinglibs-${YALANTINGLIBS_VERSION}.zip"
echo "Downloading yalantinglibs ${YALANTINGLIBS_VERSION} from ${GITHUB_PROXY}/alibaba/yalantinglibs/archive/refs/tags/${YALANTINGLIBS_VERSION}.zip"
wget -q --show-progress -O ${YALANTINGLIBS_ZIPFILE} ${GITHUB_PROXY}/alibaba/yalantinglibs/archive/refs/tags/${YALANTINGLIBS_VERSION}.zip
check_success "Failed to download yalantinglibs"
# Extract yalantinglibs
echo "Extracting yalantinglibs..."
unzip -q ${YALANTINGLIBS_ZIPFILE}
check_success "Failed to extract yalantinglibs"
# Clean up downloaded ZIP file
rm -f ${YALANTINGLIBS_ZIPFILE}
check_success "Failed to clean up downloaded ZIP file"
# Build and install yalantinglibs
cd yalantinglibs-${YALANTINGLIBS_VERSION}
check_success "Failed to change to yalantinglibs directory"
mkdir -p build
check_success "Failed to create build directory"
cd build
check_success "Failed to change to build directory"
@ -180,7 +195,171 @@ cmake --install .
check_success "Failed to install yalantinglibs"
print_success "yalantinglibs installed successfully"
cd "${REPO_ROOT}"
# Install cppzmq (From source, similar to yalantinglibs)
print_section "Installing cppzmq from source"
CPPZMQ_DIR="${THIRDPARTIES_DIR}/cppzmq"
# Check if cppzmq directory already exists
if [ -d "$CPPZMQ_DIR" ]; then
echo -e "${YELLOW}cppzmq directory already exists. Removing for fresh install...${NC}"
rm -rf "$CPPZMQ_DIR"
check_success "Failed to remove existing cppzmq directory"
fi
# Clone cppzmq
echo "Cloning cppzmq from ${GITHUB_PROXY}/zeromq/cppzmq.git"
git clone ${GITHUB_PROXY}/zeromq/cppzmq.git "$CPPZMQ_DIR"
check_success "Failed to clone cppzmq"
# Build and install cppzmq
cd "$CPPZMQ_DIR"
check_success "Failed to change to cppzmq directory"
# Checkout a specific stable version v4.11.0
echo "Checking out cppzmq version v4.11.0..."
git checkout v4.11.0
check_success "Failed to checkout cppzmq version v4.11.0"
mkdir -p build
check_success "Failed to create build directory"
cd build
check_success "Failed to change to build directory"
echo "Configuring cppzmq..."
# Key configuration: Ensure it finds the system-installed libzmq
# and configure cppzmq to install to system directory /usr/local
cmake .. \
-DCMAKE_INSTALL_PREFIX=/usr/local \
-DCPPZMQ_BUILD_TESTS=OFF
check_success "Failed to configure cppzmq"
echo "Building cppzmq (using $(nproc) cores)..."
cmake --build . -j$(nproc)
check_success "Failed to build cppzmq"
echo "Installing cppzmq..."
cmake --install .
check_success "Failed to install cppzmq"
ldconfig
print_success "cppzmq installed successfully"
# Install msgpack-cxx from release source
print_section "Installing msgpack-cxx from release source (v7.0.0)"
# Change to thirdparties directory to ensure consistent starting point
cd "$THIRDPARTIES_DIR"
check_success "Failed to change to thirdparties directory"
MSGPACK_DIR="${THIRDPARTIES_DIR}/msgpack-cxx"
# Check if msgpack directory already exists
if [ -d "$MSGPACK_DIR" ]; then
echo -e "${YELLOW}msgpack-cxx directory already exists. Removing for fresh install...${NC}"
rm -rf "$MSGPACK_DIR"
check_success "Failed to remove existing msgpack-cxx directory"
fi
# Create msgpack directory
mkdir -p "$MSGPACK_DIR"
check_success "Failed to create msgpack-cxx directory"
cd "$MSGPACK_DIR"
check_success "Failed to change to msgpack-cxx directory"
# Download and extract msgpack-cxx 7.0.0 release
MSGPACK_VERSION="7.0.0"
MSGPACK_TARBALL="msgpack-cxx-${MSGPACK_VERSION}.tar.gz"
MSGPACK_URL="${GITHUB_PROXY}/msgpack/msgpack-c/releases/download/cpp-${MSGPACK_VERSION}/${MSGPACK_TARBALL}"
echo "Downloading msgpack-cxx v${MSGPACK_VERSION} from ${MSGPACK_URL}"
wget --show-progress -O "$MSGPACK_TARBALL" "$MSGPACK_URL"
check_success "Failed to download msgpack-cxx release"
# Verify the downloaded file is not empty
if [ ! -s "$MSGPACK_TARBALL" ]; then
print_error "Downloaded msgpack-cxx tarball is empty"
fi
# Extract the tarball
echo "Extracting msgpack-cxx source..."
tar -xzf "$MSGPACK_TARBALL"
check_success "Failed to extract msgpack-cxx tarball"
# The extracted directory name includes the version
EXTRACTED_DIR="msgpack-cxx-${MSGPACK_VERSION}"
if [ ! -d "$EXTRACTED_DIR" ]; then
print_error "Extracted directory '$EXTRACTED_DIR' not found"
fi
# Move into the extracted directory
cd "$EXTRACTED_DIR"
check_success "Failed to change to extracted msgpack-cxx directory"
# Build and install msgpack-cxx
mkdir -p build
check_success "Failed to create build directory"
cd build
check_success "Failed to change to build directory"
echo "Configuring msgpack-cxx..."
cmake .. \
-DCMAKE_INSTALL_PREFIX=/usr/local \
-DMSGPACK_BUILD_TESTS=OFF \
-DMSGPACK_BUILD_EXAMPLES=OFF \
-DMSGPACK_USE_BOOST=OFF
check_success "Failed to configure msgpack-cxx"
echo "Building msgpack-cxx (using $(nproc) cores)..."
cmake --build . -j$(nproc)
check_success "Failed to build msgpack-cxx"
echo "Installing msgpack-cxx..."
cmake --install .
check_success "Failed to install msgpack-cxx"
ldconfig # Update library cache
# Clean up: remove the downloaded tarball
rm -f "../${MSGPACK_TARBALL}"
print_success "Cleaned up downloaded tarball"
print_success "msgpack-cxx v${MSGPACK_VERSION} installed successfully"
# Initialize and update git submodules
print_section "Initializing Git Submodules"
# Check if .gitmodules exists
if [ -f "${REPO_ROOT}/.gitmodules" ]; then
# Check if submodules are already initialized by looking for the .git directory in the first submodule
FIRST_SUBMODULE=$(grep "path" ${REPO_ROOT}/.gitmodules | head -1 | awk '{print $3}')
echo "Enter repository root: ${REPO_ROOT}"
cd "${REPO_ROOT}"
check_success "Failed to change to repository root directory"
if [ -d "${REPO_ROOT}/${FIRST_SUBMODULE}/.git" ] || [ -f "${REPO_ROOT}/${FIRST_SUBMODULE}/.git" ]; then
echo -e "${YELLOW}Git submodules already initialized. Skipping...${NC}"
else
echo "Initializing git submodules..."
git submodule update --init
check_success "Failed to initialize git submodules"
print_success "Git submodules initialized and updated successfully"
fi
else
echo -e "${YELLOW}No .gitmodules file found. Skipping...${NC}"
exit 1
fi
print_section "Verifying essential build tools"
@ -197,8 +376,6 @@ print_success "ldd found: $(ldd --version 2>&1 | head -1)"
print_section "Installing Go $GOVER"
USED_CN_MIRROR=false
install_go() {
ARCH=$(uname -m)
if [ "$ARCH" = "aarch64" ]; then
@ -209,45 +386,18 @@ install_go() {
echo "Unsupported architecture: $ARCH"
exit 1
fi
GO_TARBALL="go$GOVER.linux-$ARCH.tar.gz"
# Try multiple download mirrors with fallback
GO_DOWNLOAD_URLS=(
"https://go.dev/dl/${GO_TARBALL}"
"https://golang.google.cn/dl/${GO_TARBALL}"
"https://mirrors.aliyun.com/golang/${GO_TARBALL}"
)
DOWNLOAD_SUCCESS=false
for url in "${GO_DOWNLOAD_URLS[@]}"; do
echo "Downloading Go $GOVER from ${url}..."
if wget -q --show-progress --timeout=30 --tries=2 -O "${GO_TARBALL}" "${url}"; then
DOWNLOAD_SUCCESS=true
# If the official source (go.dev) failed and we fell back to a CN mirror,
# it likely means the network has restricted access to international sites.
if [[ "$url" != "https://go.dev/dl/${GO_TARBALL}" ]]; then
USED_CN_MIRROR=true
fi
print_success "Downloaded Go $GOVER from ${url}"
break
else
echo -e "${YELLOW}Failed to download from ${url}, trying next mirror...${NC}"
rm -f "${GO_TARBALL}"
fi
done
if [ "$DOWNLOAD_SUCCESS" = false ]; then
print_error "Failed to download Go $GOVER from all mirrors"
fi
# Download Go
echo "Downloading Go $GOVER..."
wget -q --show-progress https://go.dev/dl/go$GOVER.linux-$ARCH.tar.gz
check_success "Failed to download Go $GOVER"
# Install Go
echo "Installing Go $GOVER..."
tar -C /usr/local -xzf "${GO_TARBALL}"
tar -C /usr/local -xzf go$GOVER.linux-$ARCH.tar.gz
check_success "Failed to install Go $GOVER"
# Clean up downloaded file
rm -f "${GO_TARBALL}"
rm -f go$GOVER.linux-$ARCH.tar.gz
check_success "Failed to clean up Go installation file"
print_success "Go $GOVER installed successfully"
@ -273,20 +423,6 @@ if ! grep -q "export PATH=\$PATH:/usr/local/go/bin" ~/.bashrc; then
echo -e "${YELLOW}Please run 'source ~/.bashrc' or start a new terminal to use Go${NC}"
fi
# Set GOPROXY only if Go download fell back to a CN mirror, indicating restricted
# network access to international sites. Skip if user already configured GOPROXY.
if [ "$USED_CN_MIRROR" = true ] && [ -z "$GOPROXY" ]; then
export GOPROXY=https://goproxy.cn,https://goproxy.io,direct
echo -e "${YELLOW}Detected restricted network (Go was downloaded from a CN mirror).${NC}"
echo -e "${YELLOW}GOPROXY set to: ${GOPROXY}${NC}"
if ! grep -q "export GOPROXY=" ~/.bashrc; then
echo 'export GOPROXY=https://goproxy.cn,https://goproxy.io,direct' >> ~/.bashrc
echo -e "${YELLOW}GOPROXY added to ~/.bashrc for future sessions${NC}"
fi
elif [ -n "$GOPROXY" ]; then
echo -e "${GREEN}GOPROXY already set to: ${GOPROXY}${NC}"
fi
# Return to the repository root
cd "${REPO_ROOT}"

View File

@ -12,7 +12,6 @@ ENV DEBIAN_FRONTEND=noninteractive \
PYTHONUNBUFFERED=1
ARG PYTHON_VERSION=3.10
ARG PYPA_INDEX_URL=https://bootstrap.pypa.io
ARG CMAKE_BUILD_TYPE=Release
ARG EP_TORCH_VERSIONS="2.9.1"
ARG TORCH_CUDA_ARCH_LIST="8.0;9.0"
@ -23,25 +22,17 @@ ENV PYTHON_VERSION=${PYTHON_VERSION} \
TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} \
PATH="/usr/local/go/bin:${PATH}"
# Install base build utilities and the requested Python version via deadsnakes PPA
# Install base build utilities and python bindings
RUN apt-get update && \
apt-get install -y --no-install-recommends \
ca-certificates \
curl \
git \
ninja-build \
software-properties-common \
python3 \
python3-dev \
python3-pip \
python-is-python3 \
pkg-config && \
add-apt-repository -y ppa:deadsnakes/ppa && \
apt-get update && \
apt-get install -y --no-install-recommends \
python${PYTHON_VERSION} \
python${PYTHON_VERSION}-dev \
python${PYTHON_VERSION}-venv && \
curl -sS ${PYPA_INDEX_URL}/get-pip.py | python${PYTHON_VERSION} && \
update-alternatives --install /usr/bin/python python /usr/bin/python${PYTHON_VERSION} 1 && \
update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 && \
apt-get purge -y --auto-remove software-properties-common && \
rm -rf /var/lib/apt/lists/*
WORKDIR /workspace
@ -53,17 +44,16 @@ RUN bash dependencies.sh -y
# Configure & build Mooncake
RUN mkdir -p build && \
cd build && \
cmake -G Ninja .. \
cmake .. \
-DBUILD_UNIT_TESTS=OFF \
-DUSE_HTTP=ON \
-DUSE_ETCD=ON \
-DUSE_CUDA=ON \
-DWITH_EP=ON \
-DSTORE_USE_ETCD=ON \
-DPython3_EXECUTABLE=/usr/bin/python${PYTHON_VERSION} \
-DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE} && \
export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH && \
cmake --build .
cmake --build . -j"$(nproc)"
# Build nvlink allocator to make wheel self-contained for CUDA paths
RUN export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH && \
@ -85,17 +75,11 @@ ENV DEBIAN_FRONTEND=noninteractive \
PYTHONUNBUFFERED=1 \
PIP_NO_CACHE_DIR=1
# Inherit build-args so the runtime stage installs the matching interpreter
ARG PYTHON_VERSION=3.10
ARG PYPA_INDEX_URL=https://bootstrap.pypa.io
ENV PYTHON_VERSION=${PYTHON_VERSION}
# Install runtime dependencies and the requested Python version
# Install runtime dependencies required by Mooncake
RUN apt-get update && \
apt-get install -y --no-install-recommends \
ca-certificates \
curl \
software-properties-common \
python3 \
python3-pip \
ibverbs-providers \
rdma-core \
libibverbs1 \
@ -104,18 +88,10 @@ RUN apt-get update && \
liburing2 \
libyaml-0-2 \
libcurl4 && \
add-apt-repository -y ppa:deadsnakes/ppa && \
apt-get update && \
apt-get install -y --no-install-recommends \
python${PYTHON_VERSION} && \
curl -sS ${PYPA_INDEX_URL}/get-pip.py | python${PYTHON_VERSION} && \
update-alternatives --install /usr/bin/python python /usr/bin/python${PYTHON_VERSION} 1 && \
update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 && \
apt-get purge -y --auto-remove software-properties-common curl && \
rm -rf /var/lib/apt/lists/*
# Copy wheels produced in builder stage and install them via pip
COPY --from=builder /workspace/mooncake-wheel/dist /tmp/mooncake-wheel
RUN python${PYTHON_VERSION} -m pip install --no-cache-dir /tmp/mooncake-wheel/*.whl && rm -rf /tmp/mooncake-wheel /root/.cache/pip
RUN python3 -m pip install --no-cache-dir /tmp/mooncake-wheel/*.whl && rm -rf /tmp/mooncake-wheel /root/.cache/pip
CMD ["/bin/bash"]

View File

@ -8,7 +8,6 @@ This page summarizes useful flags, environment variables, and HTTP endpoints to
- `--rpc_port` (int, default 50051): RPC listen port.
- `--rpc_thread_num` (int, default min(4, CPU cores)): RPC worker threads. If not set, uses `--max_threads` (default 4) capped by CPU cores.
- `--rpc_address` (str, default `0.0.0.0`): RPC bind address.
- `--rpc_interface` (str, default empty): Network interface used to resolve the final RPC address. When set, Mooncake Master resolves the interface's current IPv4 address at startup and uses it as the final `rpc_address`. This overrides `--rpc_address`.
- `--rpc_conn_timeout_seconds` (int, default `0`): RPC idle connection timeout; `0` disables.
- `--rpc_enable_tcp_no_delay` (bool, default `true`): Enable TCP_NODELAY.
@ -27,8 +26,8 @@ This page summarizes useful flags, environment variables, and HTTP endpoints to
- `free_ratio_first`: Free-ratio-first strategy. Samples multiple candidates and selects those with highest free space ratio for better load balancing.
- Eviction and TTLs
- `--default_kv_lease_ttl` (duration, default `5000` ms): Default lease TTL for KV objects. The default unit is milliseconds, so `5000` means `5000ms`. Duration strings such as `5000ms`, `5s`, `30m`, or `1h` are also supported.
- `--default_kv_soft_pin_ttl` (duration, default `1800000` ms): Soft pin TTL (30 minutes). The default unit is milliseconds, so `1800000` means `1800000ms`. Duration strings such as `1800000ms`, `30m`, or `1h` are also supported.
- `--default_kv_lease_ttl` (uint64, default `5000` ms): Default lease TTL for KV objects.
- `--default_kv_soft_pin_ttl` (uint64, default `1800000` ms): Soft pin TTL (30 minutes).
- `--allow_evict_soft_pinned_objects` (bool, default `true`): Allow evicting soft-pinned objects.
- `--eviction_ratio` (double, default `0.05`): Fraction evicted when hitting high watermark.
- `--eviction_high_watermark_ratio` (double, default `0.95`): Usage ratio to trigger eviction.
@ -75,18 +74,6 @@ mooncake_master \
--enable_metric_reporting=true
```
Example (resolve the master RPC address from a stable interface name in a container):
```bash
mooncake_master \
--rpc_interface=eth0 \
--enable_http_metadata_server=true \
--http_metadata_server_host=0.0.0.0 \
--http_metadata_server_port=8080
```
This resolves the current IPv4 address of `eth0` at startup and uses it as the final `rpc_address`.
Example (use free-ratio-first allocation strategy for better load balancing):
```bash
@ -105,13 +92,6 @@ mooncake_master \
--config_path=mooncake-store/conf/master.yaml
```
For config files, the equivalent setting is:
```yaml
rpc_interface: "eth0"
rpc_port: 50051
```
## Metrics Endpoints
The master exposes Prometheus-style metrics over HTTP on `--metrics_port`:
@ -158,13 +138,3 @@ Available log levels: trace, debug, info, warn (or warning), error, and critical
- Scale `--rpc_thread_num` with available CPU cores and workload.
- Start with default eviction settings; adjust `--eviction_high_watermark_ratio` and `--eviction_ratio` based on memory pressure and object churn.
- Use `/metrics/summary` during bring-up; integrate `/metrics` with Prometheus/Grafana for production.
---
:::{toctree}
:caption: Advanced Topics
:maxdepth: 1
ssd-offload
:::

View File

@ -1,280 +0,0 @@
# SSD Offload
## Overview
Mooncake Store supports offloading KV cache objects from distributed memory to local SSD. When memory pressure is high, the master instructs clients to persist selected objects to disk. On a cache miss, the client automatically falls back to reading from SSD.
SSD offload is currently **only available in Real Client mode**. The real client is a standalone process that communicates with the application (e.g., SGLang) via RPC. All SSD reads and writes happen within this process.
## Startup Steps
### Step 1: Create the SSD storage directory
```bash
mkdir -p /nvme/mooncake_offload
```
### Step 2: Start the master
```bash
mooncake_master \
--rpc_port=50051 \
--enable-offload true
```
### Step 3: Start the real client with SSD offload enabled
Use the `--enable_offload` flag to enable SSD offload, and set environment variables to specify the storage path and backend:
```bash
export MOONCAKE_OFFLOAD_FILE_STORAGE_PATH=/nvme/mooncake_offload
export MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR=bucket_storage_backend
mooncake_client \
--master_server_address=127.0.0.1:50051 \
--host=<machine IP> \
--protocol="rdma" \
--device_names=<NIC name, e.g. eth0> \
--port=50052 \
--global_segment_size="4 GB" \
--enable_offload=true \
--metadata_server="P2PHANDSHAKE"
```
> **Note:** On startup, the real client automatically scans existing SSD data and reports it to the master. No manual recovery is needed.
### Step 4: Connect the application to the real client
The application (e.g., SGLang) connects to the real client via the `MooncakeDistributedStore` Python SDK. SSD offload and fallback loading are handled transparently.
```python
from mooncake.store import MooncakeDistributedStore
store = MooncakeDistributedStore()
store.setup(
local_hostname="<machine IP>",
metadata_server="P2PHANDSHAKE",
global_segment_size=4 * 1024 * 1024 * 1024, # 4 GB
local_buffer_size=512 * 1024 * 1024, #512MB
protocol="rdma",
device_name="eth0",
master_server_address="127.0.0.1:50051",
)
```
---
## Real Client Parameters
| Flag | Default | Description |
|------|---------|-------------|
| `--master_server_address` | `127.0.0.1:50051` | Master address |
| `--host` | `0.0.0.0` | This machine's externally reachable IP |
| `--port` | `50052` | Real client RPC listening port |
| `--device_names` | ` ` | NIC name(s), e.g. `eth0` or `mlx5_0` |
| `--protocol` | `tcp` | Transport protocol: `tcp` or `rdma` |
| `--global_segment_size` | `4 GB` | Memory pool size allocated for this node |
| `--enable_offload` | `false` | **Must be set to `true` to enable SSD offload** |
| `--threads` | `1` | Number of RPC server threads |
---
## SSD Offload Configuration
### Core settings
| Environment Variable | Default | Description |
|---|---|---|
| `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 |
### Bucket backend settings
Applies when `MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR=bucket_storage_backend`.
| Environment Variable | Default | Description |
|---|---|---|
| `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` |
---
## Storage Backends
### `bucket_storage_backend` (recommended)
Groups multiple objects into bucket files. Reduces filesystem overhead, supports efficient batch I/O, and supports FIFO and LRU eviction.
**File layout:**
```
/nvme/mooncake_offload/
├── 1710000000000-0.bucket # data file (multiple KV pairs)
├── 1710000000000-0.meta # metadata file
├── 1710000000001-0.bucket
└── ...
```
Best for: general-purpose use, large-scale deployments.
### `file_per_key_storage_backend`
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`
Pre-allocates a single large file and manages offset-based allocation within it. Highest concurrency via 1024-shard metadata.
> **Warning:** This backend does **not** support metadata recovery on restart. On initialization, the data file is truncated and all in-memory metadata is cleared. Any previously offloaded objects become inaccessible after a process restart.
**Capacity:** `MOONCAKE_OFFLOAD_TOTAL_SIZE_LIMIT_BYTES` is used directly as the pre-allocated file size (100%, no safety margin). Unlike `bucket_storage_backend`, there is no separate quota variable — this is the sole disk usage control. Set it below the physical disk capacity to avoid filling the disk; writes are rejected once usage reaches this limit.
Best for: high-concurrency scenarios with many small objects where restart durability is not required.
---
## 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.
| Policy | Behavior |
|--------|----------|
| `none` | No eviction (default); writes fail when disk is full |
| `fifo` | Evict the oldest bucket first |
| `lru` | Evict the least recently read bucket first |
Eviction is two-phase: the bucket is removed from metadata and master is notified first, then in-flight reads are drained before files are deleted.
---
## Example
The following example starts a master and a real client on a single machine.
### Environment
- Machine IP: `192.168.1.10`
- NIC: `eth0`
- SSD mount point: `/nvme`
- Memory pool size: 4 GB (smaller than the total data written, to trigger offload)
### Start the master
```bash
mooncake_master \
--rpc_port=50051
```
### Start the real client (new terminal)
```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
mooncake_client \
--master_server_address="192.168.1.10:50051" \
--host="192.168.1.10" \
--device_names="eth0" \
--port=50052 \
--protocol="rdma" \
--global_segment_size="4GB" \
--enable_offload="true"
```
---
## Notes
- `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` must be an absolute path to an existing, writable directory. Symbolic links and paths containing `..` are rejected.
- On real client restart, the backend automatically scans existing SSD files and reports them to the master, so previously offloaded objects remain accessible.
- Eviction only notifies the master and deletes local files; objects replicated on other nodes are unaffected.
- Each machine requires its own real client process. In multi-node deployments, ensure `--host` and `--port` are correctly set so nodes can reach each other.
**2-node example:** suppose Node A (`192.168.1.10`) runs the master and Node B (`192.168.1.11`) is a second worker. Both real clients must point to the same master and advertise their own externally reachable IP:
```bash
# Node A — runs the master and its own real client
mooncake_master --rpc_port=50051 --enable-offload true &
export MOONCAKE_OFFLOAD_FILE_STORAGE_PATH=/nvme/mooncake_offload
mooncake_client \
--master_server_address="192.168.1.10:50051" \
--host="192.168.1.10" \ # externally reachable IP of Node A
--device_names="eth0" \
--protocol="rdma" \
--metadata_server="P2PHANDSHAKE" \
--port=50052 \
--global_segment_size="4GB" \
--enable_offload="true"
```
```bash
# Node B — real client only; points to the same master on Node A
export MOONCAKE_OFFLOAD_FILE_STORAGE_PATH=/nvme/mooncake_offload
mooncake_client \
--master_server_address="192.168.1.10:50051" \
--host="192.168.1.11" \ # externally reachable IP of Node B, NOT 127.0.0.1
--device_names="eth0" \
--protocol="rdma" \
--metadata_server="P2PHANDSHAKE" \
--port=50052 \
--global_segment_size="4GB" \
--enable_offload="true"
```
---
## Troubleshooting
### SSD offload is not triggering
- Confirm `--enable_offload=true` is passed to `mooncake_client` and `--enable-offload true` is passed to `mooncake_master`.
- Check that `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` points to an existing, writable directory. The client will fail silently if the path is invalid.
- Verify memory pressure is actually high enough for the master to trigger offload. If the memory pool (`--global_segment_size`) is large relative to the data written, offload may never activate.
### "Permission denied" or "No such file or directory" on the storage path
- Ensure the directory exists before starting the client: `mkdir -p <path>`.
- Confirm the process user has read/write access to the directory.
- Symbolic links and paths containing `..` are rejected — use an absolute, canonical path.
### "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.
Check the current limit:
```bash
ulimit -l # in KB; "unlimited" means no cap
```
To raise it for the current session:
```bash
ulimit -l unlimited
```
To raise it permanently, add the following to `/etc/security/limits.conf`:
```
* soft memlock unlimited
* hard memlock unlimited
```
Alternatively, reduce `MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES` to a value within the existing limit. Note that the warning does not abort startup — the client falls back to non-fixed-buffer I/O — but performance may be lower than expected.

View File

@ -0,0 +1,430 @@
# KV Event Subscriber Guide
This document provides guidance for developers who wish to subscribe to KV cache events in the Mooncake system. It focuses on the message schema and deserialization methods needed to properly handle events sent by the event publisher.
---
## Event Types & Meanings
Before diving into the technical schema, it's important to understand what each event type represents and when they are triggered in the Mooncake system.
### Event Overview Table
| Event Type | Trigger Condition | Purpose | Key Characteristics |
|------------|-------------------|---------|---------------------|
| **BlockStoreEvent** | First storage of KV cache block | Initial block creation with metadata | ✅ contains `StoreEventInfo` metadata<br>✅ Represents initial creation<br>✅ Generated during `Put` operations |
| **BlockUpdateEvent** | Replica management operations (copy/remove/migrate) | Track replica distribution changes | ❌ No `StoreEventInfo` metadata<br>✅ Focuses on replica locations<br>✅ Internal system operations |
| **RemoveAllEvent** | System-wide cache clearance | Signal cache invalidation | ❌ No additional fields<br>✅ System-wide invalidation<br>✅ Maintenance operations |
### Event Type Quick Reference
#### 🏗️ BlockStoreEvent - Storage Event
```mermaid
graph LR
A[Put Operation] --> B[BlockStoreEvent]
B --> C[Contains Metadata]
C --> D[Downstream Processing]
```
#### 🔄 BlockUpdateEvent - Replica Event
```mermaid
graph LR
A[Replica Operation] --> B[BlockUpdateEvent]
B --> C[Replica Location Changes]
C --> D[Downstream Processing]
```
#### 🧹 RemoveAllEvent - Clearance Event
```mermaid
graph LR
A[System Maintenance] --> B[RemoveAllEvent]
B --> C[Cache Invalidation]
C --> D[Downstream Processing]
```
### Field Definitions at a Glance
#### Core Event Fields
| Field | Type | Required | Description |
|-------|------|----------|-------------|
| `mooncake_key` | `std::string` | ✅ All events | Unique identifier for cached object |
| `replicas` | Nested array | ✅ BlockStore/Update | Replica locations: `[type, location]` |
#### `StoreEventInfo` Fields (`BlockStoreEvent` only)
When processing `BlockStoreEvent`events, pay special attention to the StoreEventInfo fields which are appended to the event. **These fields have specific default values that indicate when they are not set**:
| Field | Type | Default | Description |
|-------|------|---------|-------------|
| `model_name` | `std::string` | `""` | Model identifier |
| `block_size` | `uint32_t` | `0` | Block size in bytes |
| `block_hash` | `std::string` | `""` | Current block hash |
| `parent_block_hash` | `std::string` | `""` | Parent block hash |
| `token_ids` | `std::vector<uint32_t>` | `[]` | Token ID sequence |
---
## Event Message Schema
The KV event system uses ZeroMQ (ZMQ) for message transport and MessagePack for serialization. Events are sent as multipart ZMQ messages containing three components:
### ZMQ Message Structure
Each event message consists of 3 ZMQ message parts:
```plainText
1. topic Part: Contains the topic string (default: "mooncake")
2. Sequence Number Part: 8-byte big-endian unsigned integer representing the sequence number
3. Payload Part: MessagePack-serialized event batch data
```
### Event Batch Schema
The payload part contains a serialized `EventBatch`object with the following structure:
```typescript
[
ts, // First element: Timestamp (double) - seconds since UNIX epoch
[ // Second element: Event list (array)
[event1_data], // Event 1
[event2_data], // Event 2
... // Additional events
]
]
```
Each event in the batch is a serialized event object that begins with an event type identifier string followed by its specific data fields.
### Event Types Schema
There are three main event types supported by the system:
#### Schema
```typescript
// Event triggered on the first storage occurrence - contains 8 fields
BlockStoreEvent {
"BlockStoreEvent", // Event type identifier, string type
std::string mooncake_key, // Mooncake key
[ // Replica location list (nested arrays)
["memory", "transport_endpoint"], // Memory replica
["disk", "file_path"], // Disk replica
["local_disk", "transport_endpoint"] // Local disk replica
],
std::string model_name, // Model name
uint32_t block_size, // Block size
std::string block_hash, // Current block hash
std::string parent_block_hash, // Parent block hash
std::vector<uint32_t> token_ids // Token ID sequence
}
// Contains all internal Mooncake system operations - contains 3 fields
BlockUpdateEvent {
"BlockUpdateEvent", // Event type identifier
std::string mooncake_key,
[ // Replica location list
["memory", "transport_endpoint"],
["disk", "file_path"],
["local_disk", "transport_endpoint"]
]
}
// Event for clearing all KV Cache in Mooncake - contains 1 field
RemoveAllEvent {
"RemoveAllEvent" // Event type identifier
}
```
#### Example
```typescript
// BlockStoreEvent
[
"BlockStoreEvent", // Event type identifier
"key_12345", // mooncake key
[ // Replica location list (nested arrays)
["memory", "tcp://192.168.1.10:6000"], // Memory type location
["disk", "/data/blocks/block_12345.bin"], // Disk type location
["local_disk", "tcp://192.168.1.10:7000"] // Local disk type location
],
"llama2-7b", // Model name
512, // block size
"0x41234125", // Current block hash
"0x51512342", // Parent block hash
[1,2,3,4,5] // Token id list
]
// BlockUpdateEvent
[
"BlockUpdateEvent", // Event type identifier
"key_12345", // mooncake key
[ // Replica location list (nested arrays)
["memory", "tcp://192.168.1.10:6000"], // Memory type location
],
]
// RemoveAllEvent
[
"RemoveAllEvent" // Event type identifier
]
```
---
## Deserialization Steps
To properly deserialize events from the KV event system, follow these steps:
1. **Receive the multipart message**: Use `zmq::recv_multipart`or equivalent to receive all 3 parts of the message.
2. **Extract the topic**: The first part contains the topic string.
3. **Extract and convert the sequence number**: The second part contains an 8-byte big-endian unsigned integer. On little-endian systems, convert it using `be64toh()`or equivalent function.
4. **Deserialize the payload**: The third part contains MessagePack-serialized data. Deserialize it to get the `EventBatch`.
5. **Process individual events**: Iterate through the events in the batch and handle each according to its type identifier (the first element in each event array).
## Required Libraries
- **ZeroMQ library**: For receiving multipart messages
- **MessagePack library**: For deserializing the payload
- **Byte order conversion functions**: For converting sequence numbers between big-endian and host byte order
---
## Event Batch Timestamp Handling
The event batch timestamp field `ts` is a double-precision floating-point number representing seconds since the UNIX epoch (January 1, 1970). When deserializing, convert this timestamp appropriately:
<details>
<summary>Click to expand: Python example</summary>
```python
import datetime
# Python example
timestamp = event_batch[0] # double type seconds
dt = datetime.datetime.fromtimestamp(timestamp)
print(f"Event batch timestamp: {dt}")
```
</details>
<details>
<summary>Click to expand: GoLang example</summary>
```go
// Go example
timestamp := eventBatch[0].(float64)
t := time.Unix(int64(timestamp), 0)
fmt.Printf("Event batch timestamp: %v\n", t)
```
</details>
---
## Python Example
Here's a Python example showing how to subscribe to and deserialize KV events:
<details>
<summary>Click to expand: Python example</summary>
```python
def deserialize_block_store_event(event_array):
if len(event_array) < 8:
raise ValueError("Invalid BlockStoreEvent array length")
# Extract fields by fixed position
event_type = event_array[0] # "BlockStoreEvent"
mooncake_key = event_array[1]
replicas = event_array[2] # Nested array of replica locations
model_name = event_array[3] # String, defaults to ""
block_size = event_array[4] # Integer, defaults to 0
block_hash = event_array[5] # String, defaults to ""
parent_block_hash = event_array[6] # String, defaults to ""
token_ids = event_array[7] # List of integers, defaults to []
# Handle default values according to StoreEventInfo specifications
if model_name == "":
# Model name not set, use default or skip processing
model_name = "unknown"
if block_size == 0:
# Invalid block size, may indicate an error
raise ValueError("Invalid block size: 0")
if block_hash == "":
# block_hash is not set, use None to indicate missing value
block_hash = None
if parent_block_hash == "":
# No parent block (root block), handle appropriately
parent_block_hash = None
return {
"type": event_type,
"mooncake_key": mooncake_key,
"replicas": replicas,
"model_name": model_name,
"block_size": block_size,
"block_hash": block_hash,
"parent_block_hash": parent_block_hash,
"token_ids": token_ids
}
def process_replica_locations(replicas):
"""Process nested replica location arrays"""
replica_info = []
for replica in replicas:
if len(replica) != 2:
continue
replica_type = replica[0] # "memory", "disk", or "local_disk"
location = replica[1] # Endpoint or file path
replica_info.append({"type": replica_type, "location": location})
return replica_info
```
</details>
## GoLang Example
Here's a Go example showing how to subscribe to and deserialize KV events:
<details>
<summary>Click to expand: GoLang example</summary>
```go
func deserializeBlockStoreEvent(eventSlice []interface{}) (map[string]interface{}, error) {
if len(eventSlice) < 8 {
return nil, fmt.Errorf("invalid BlockStoreEvent array length: %d", len(eventSlice))
}
result := make(map[string]interface{})
// Extract fields by fixed position
result["type"] = eventSlice[0].(string)
result["mooncake_key"] = eventSlice[1].(string)
result["replicas"] = eventSlice[2]
result["model_name"] = eventSlice[3].(string)
// Handle block_size with proper type assertion
if blockSize, ok := eventSlice[4].(uint32); ok {
result["block_size"] = blockSize
if blockSize == 0 {
return nil, fmt.Errorf("invalid block size: 0")
}
} else {
return nil, fmt.Errorf("invalid block_size type")
}
// block_hash field handling with default value check
blockHash := eventSlice[5].(string)
if blockHash == "" {
result["block_hash"] = nil // Indicates missing value
} else {
result["block_hash"] = blockHash
}
// parent_block_hash field handling
parentBlockHash := eventSlice[6].(string)
if parentBlockHash == "" {
result["parent_block_hash"] = nil // No parent block
} else {
result["parent_block_hash"] = parentBlockHash
}
result["token_ids"] = eventSlice[7]
return result, nil
}
func processReplicas(replicas interface{}) ([]map[string]string, error) {
// Process nested replica arrays
replicaSlice, ok := replicas.([]interface{})
if !ok {
return nil, fmt.Errorf("invalid replicas type")
}
var replicaInfo []map[string]string
for _, replica := range replicaSlice {
replicaArr, ok := replica.([]interface{})
if !ok || len(replicaArr) != 2 {
continue
}
replicaType, ok1 := replicaArr[0].(string)
location, ok2 := replicaArr[1].(string)
if ok1 && ok2 {
replicaInfo = append(replicaInfo, map[string]string{
"type": replicaType,
"location": location,
})
}
}
return replicaInfo, nil
}
```
</details>
---
## Replay Functionality
The KV event system supports replay functionality, allowing subscribers to request historical events. To implement replay:
1. **Send replay request**: Send a 3-part ZMQ message to the replay endpoint:
- Part 1: Client identifier (any data that identifies your client)
- Part 2: Empty frame
- Part 3: Starting sequence number (8-byte big-endian unsigned integer)
2. **Receive replay events**: The system will send historical events starting from the requested sequence number. Each replayed event is sent as a regular ZMQ multipart message.
3. **Replay end marker**: When all available historical events have been sent, the system sends a special end marker message containing the magic sequence `0xFFFFFFFFFFFFFFFF` (8 bytes of 0xFF).
4. **Handle replay completion**: Upon receiving the end marker, you know that the replay session has completed.
5. **Important considerations**:
- Replayed events may be older than real-time events already processed
- Implement duplicate detection if needed
- Replay buffer size is limited by the system configuration
Example replay request handling:
<details>
<summary>Click to expand: Python example</summary>
```python
def send_replay_request(socket, start_seq):
"""Send replay request to the replay endpoint"""
client_id = b"my_client_001" # Your client identifier
empty_frame = b"" # Empty frame
seq_be = start_seq.to_bytes(8, byteorder='big') # Big-endian sequence
messages = [client_id, empty_frame, seq_be]
socket.send_multipart(messages)
def handle_replay_message(messages):
"""Handle incoming replay message"""
if len(messages) == 4:
client_id = messages[0]
empty_frame = messages[1]
seq_be = messages[2]
payload = messages[3]
# Check for end marker
if payload == b'\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF':
print("Replay session completed")
return None
# Process regular replay event
seq = int.from_bytes(seq_be, byteorder='big')
return deserialize_payload(payload)
return None
```
</details>

View File

@ -0,0 +1,459 @@
# Mooncake Event System Developer Manual
## 1. System Overview
KVEventSystem is an asynchronous event processing system based on the publish-subscribe pattern. It encapsulates the complex logic of event publishing, consumption, and queue management using the Facade Pattern, providing a concise and unified interface for upper-layer applications.
### 1.1 Core Features
- **Asynchronous Event Publishing**: Supports publishing various types of events in a non-blocking manner.
- **Batch Processing**: Automatically batches and merges events to optimize network transmission efficiency.
- **Reliable Transmission**: Built-in retry mechanisms and queue buffering ensure no event loss.
- **Real-time Monitoring**: Provides detailed runtime statistical information.
- **Flexible Configuration**: Performance and reliability can be adjusted through configuration parameters.
### 1.2 Architecture Design
#### 1.2.1 Architecture Logic Diagram
```mermaid
graph TD
ApplicationLayer[API call site] -->|invoke interface<br>publish KVEvent| KVEventProducer
IR[indexer/router]
subgraph KVEventSystem
KVEventProducer -->|enqueue| KVEventQueue
subgraph KVEventConsumer
ZMQ
end
KVEventQueue -->|dequeue| KVEventConsumer
end
ZMQ -->|event batch serialization<br>ZMQ network transmission| IR
```
#### 1.2.2 Class Relationship Diagram
```mermaid
classDiagram
%% Event Class Hierarchy
class KVCacheEvent {
<<interface>>
+pack(msgpack::packer~msgpack::sbuffer~) void
+type_tag() string_view
}
KVCacheEvent <|-- BlockStoreEvent
KVCacheEvent <|-- BlockUpdateEvent
KVCacheEvent <|-- RemoveAllEvent
%% Core Components
class KVEventProducer {
-shared_ptr~KVEventQueue~ event_queue_
-unique_ptr~ThreadPool~ enqueue_pool_ : enqueue worker(1)
+publish~Event, Args...~(args) future~bool~
+publish_event_async(KVEventPtr) future~bool~
+shutdown() void
}
class KVEventConsumer {
-zmq::context_t context_
-shared_ptr~KVEventQueue~ event_queue_
-jthread publisher_thread_
+KVEventConsumer(shared_ptr~KVEventQueue~, Config)
+shutdown() void
+is_running() bool
}
class ThreadSafeQueue~T~ {
<<alias KVEventQueue>>
}
%% Facade Class
class KVEventSystem {
-shared_ptr~KVEventQueue~ event_queue_
+KVEventSystem(const KVEventPublisherConfig&)
+~KVEventSystem()
+shutdown() void
+is_running() bool
+publish~Event, Args...~(args) future~bool~
}
%% Configuration and Event Classes
class KVEventPublisherConfig {
}
class StoreEventInfo {
}
class EventBatch {
+serialize() msgpack::sbuffer
}
%% Relationships
KVEventSystem --> KVEventPublisherConfig : reads
KVEventSystem --> KVEventProducer : contains
KVEventSystem --> KVEventConsumer : contains
KVEventSystem --> ThreadSafeQueue~KVEventPtr~ : contains
KVEventProducer --> ThreadSafeQueue~KVEventPtr~ : enqueues
KVEventConsumer --> ThreadSafeQueue~KVEventPtr~ : dequeues
KVEventProducer --> KVCacheEvent : creates
BlockStoreEvent --> StoreEventInfo : contains
KVCacheEvent <.. EventBatch : aggregates
EventBatch <.. KVEventConsumer : serialize & publish
```
#### 1.2.3 Data Flow Sequence Diagram
```mermaid
sequenceDiagram
participant App as APICallSite
participant ES as KVEventSystem
participant EP as KVEventProducer
participant TSQ as ThreadSafeQueue
participant ZMQ as KVEventConsumer
participant Network as ZeroMQNetwork
App->>ES: 1. publish<KVEvent>(...)
ES->>EP: 2. publish<KVEvent>(...)
EP->>TSQ: 3. push(event)<br>(create KVCacheEvent object)
Note over TSQ: KVEvent Buffer
ZMQ->>TSQ: 4. peek_batch()
ZMQ->>ZMQ: 5. batch processing
ZMQ->>Network: 6. ZeroMQ Publish
Network-->>ZMQ: 7. Publish confirmation
ZMQ-->>TSQ: 8. pop_batch()
ES-->>App: 9. return future
Note over App,Network: Asynchronous processing flow
```
------
## 2. Key Class Interface Declarations
### 2.1 Event Classes (kv_event.hpp)
#### Event Class Relationships
| Class Name | Parent Class | Key Fields | Description |
| ------------------ | -------------- | ---------------------------------------------- | --------------------------------------------- |
| `KVCacheEvent` | - | No fields | Abstract base class, defines event interface |
| `StoreEventInfo` | - | All fields | Container for additional data of store events |
| `BlockStoreEvent` | `KVCacheEvent` | `mooncake_key`, `replicas`, `store_event_info` | Block storage event |
| `BlockUpdateEvent` | `KVCacheEvent` | `mooncake_key`, `replicas` | Block update event |
| `RemoveAllEvent` | `KVCacheEvent` | No fields | Clear all cache event |
**UML Class Diagram**:
```mermaid
classDiagram
%% Event Class Hierarchy
class KVCacheEvent {
<<abstract>>
+pack(msgpack::packer~msgpack::sbuffer~& pk) void
+type_tag() string_view
}
KVCacheEvent <|-- BlockStoreEvent
KVCacheEvent <|-- BlockUpdateEvent
KVCacheEvent <|-- RemoveAllEvent
%% StoreEventInfo Structure
class StoreEventInfo {
+std::string model_name
+uint32_t block_size
+std::string block_hash
+std::string parent_block_hash
+std::vector~uint32_t~ token_ids
}
%% BlockStoreEvent Class
class BlockStoreEvent {
+std::string mooncake_key
+std::vector~Replica::Descriptor~ replicas
+StoreEventInfo store_event_info
+BlockStoreEvent(key, replica_list, info)
+pack(pk) void
+type_tag() string_view
}
%% BlockUpdateEvent Class
class BlockUpdateEvent {
+std::string mooncake_key
+std::vector~Replica::Descriptor~ replicas
+BlockUpdateEvent(key, replica_list)
+pack(pk) void
+type_tag() string_view
}
%% RemoveAllEvent Class
class RemoveAllEvent {
+RemoveAllEvent()
+pack(pk) void
+type_tag() string_view
}
%% Relationships
BlockStoreEvent *-- StoreEventInfo : contains
```
------
#### Event Class Types
##### Event Base Class: `KVCacheEvent`
> The **abstract base class** for all KV cache events, defining the **unified interface** for event serialization and type identification. It uses the abstract base class pattern to provide a unified event processing interface. All concrete event types must inherit from this class and implement the interface methods.
```c++
struct KVCacheEvent {
virtual void pack(msgpack::packer<msgpack::sbuffer>& pk) const = 0; // Serialization method
virtual std::string_view type_tag() const = 0; // Type identification method
}
```
**Explanation**
| Method Name | Return Type | Parameters | Description |
| ----------- | ------------------ | --------------------------------------- | ------------------------------------------------------------ |
| `pack` | `void` | `msgpack::packer<msgpack::sbuffer>& pk` | Pure virtual function, serializes the event into MessagePack format. Derived classes must implement this method to provide type-specific serialization logic. |
| `type_tag` | `std::string_view` | None | Pure virtual function, returns a string view of the event type identifier, used for type identification during deserialization. |
------
##### Store Event Pass-Through Meta Data Structure: `StoreEventInfo`
> The `StoreEventInfo`struct is used to carry business metadata for KV cache storage events. All fields in this structure are optional; which fields to populate should be determined by the specific business scenario of the upper-layer cache-aware component (such as indexer or router). The Mooncake system acts only as a transparent pipeline for transmitting these fields and will not perform any business logic validation or interpretation of their content.
```c++
struct StoreEventInfo {
std::string model_name{""};
uint32_t block_size{0};
std::string block_hash{""};
std::string parent_block_hash{""};
std::vector token_ids{};
};
```
**Explanation**
| Field Name | Type | Default Value | Constraints/Explanation |
| ------------------- | ----------------------- | ------------- | --------------------------------------------------------- |
| `model_name` | `std::string` | `""` | Model name identifier, empty indicates not set |
| `block_size` | `uint32_t` | `0` | Block size (bytes), 0 indicates invalid value |
| `block_hash` | `std::string` | `""` | Hash of the current block, used for unique identification |
| `parent_block_hash` | `std::string` | `""` | Parent block hash, used for dependency chain |
| `token_ids` | `std::vector<uint32_t>` | Empty vector | Token ID sequence, can be empty |
------
##### Storage Event Structure: `BlockStoreEvent`
> Event type triggered when a KV cache block is stored for the first time, typically containing the additional storage information `StoreEventInfo`that needs to be passed.
```c++
struct BlockStoreEvent {
std::string mooncake_key;
std::vector<Replica::Descriptor> replicas;
StoreEventInfo store_event_info;
}
```
**Explanation**
| Field Name | Type | Default Value (set by constructor) | Constraints/Explanation |
| ------------------ | ---------------------------------- | ---------------------------------- | -------------------------------------------- |
| `mooncake_key` | `std::string` | N/A | Unique identifier of the cached `ObjectMetadata` in the `MasterService` |
| `replicas` | `std::vector<Replica::Descriptor>` | N/A | List of replica descriptors containing replica location and type information |
| `store_event_info` | `StoreEventInfo` | N/A | Additional information required for KVCache-aware algorithms |
> The `mooncake_key` field represents the primary key of the cached object in the distributed KV cache system. In the implementation, this value may originate from method parameters (e.g., the key parameter in `PutEnd` method) or from `it->first` when iterating through metadata maps (where it is an iterator over `std::unordered_map<std::string, ObjectMetadata>` in `MasterService::metadata_shards_`).
**Event Format Example**
```javascript
[
"BlockStoreEvent", // Event type identifier
"key_001", // mooncake key
[ // Replica location list
["memory", "tcp://192.168.1.10:6000"], // Memory type location
["disk", "/data/blocks/block_12345.bin"], // Disk type location
],
"llama2-7b", // Model name
512, // block size
"0x41234125", // Current block hash
"0x51512342", // Parent block hash
[1,2,3,4,5] // Token id list
]
```
------
##### Replica Change Event Structure: `BlockUpdateEvent`
> Event type triggered by any replica update operation other than storage events, including copy/eviction/migration.
```c++
struct BlockUpdateEvent {
std::string mooncake_key;
std::vector<Replica::Descriptor> replicas;
}
```
**Explanation**
| Field Name | Type | Default Value | Constraints/Explanation |
| -------------- | ---------------------------------- | ------------- | ----------------------------------------------- |
| `mooncake_key` | `std::string` | N/A | Unique identifier of the cached `ObjectMetadata` in the `MasterService` |
| `replicas` | `std::vector<Replica::Descriptor>` | N/A | List of replica descriptors containing replica location and type information |
**Event Format Example**
```javascript
[
"BlockUpdateEvent", // Event type identifier
"key_12345", // mooncake key
[ // Global replica location list
["memory", "tcp://192.168.1.10:6000"],
["memory", "tcp://192.168.1.11:5001"],
["disk", "/data/blocks/block_12345.bin"],
],
]
```
------
##### Clear All Cache Event Structure: `RemoveAllEvent`
> Event triggered to clear all caches in Mooncake-Store.
```c++
RemoveAllEvent {
}
```
**Explanation**: *This event has no fields, identified only by its type tag.*
**Event Format Example**
```javascript
[
"RemoveAllEvent"
]
```
------
#### Serialization-Related Types
##### Event Batch Structure: `EventBatch`
> Used to batch serialize events for transmission, thereby reducing network overhead and improving throughput, while timestamps ensure event ordering.
```c++
struct EventBatch {
double ts;
std::vector<std::shared_ptr<KVCacheEvent>> events;
}
```
**Explanation**
| Field Name | Type | Default Value (set by constructor) | Constraints/Explanation |
| ---------- | -------------------------------------------- | ---------------------------------- | ------------------------------------------------------------ |
| `ts` | `double` | N/A | Timestamp of batch creation, using Unix timestamp (seconds since 1970-01-01 00:00:00 UTC, floating-point), used by the receiver to process batches in chronological order. |
| `events` | `std::vector<std::shared_ptr<KVCacheEvent>>` | N/A | Array of event pointers, can contain any event object derived from `KVCacheEvent`. Events in the batch are serialized and processed in array order. |
**Event Batch Format Example**
The event batch is serialized as an array containing two elements:
- The first element is the timestamp (floating-point number)
- The second element is an array of events, where each event is its corresponding type's serialized array.
```javascript
[
170000000.123, // Element 1: Timestamp (double)
[ // Element 2: Event list (array)
[ // Event 1: BlockStoreEvent
"BlockStoreEvent",
"key_001",
[
["memory", "tcp://192.168.1.10:6000"],
["disk", "/data/blocks/block_12345.bin"]
],
"llama2-7b",
512,
"0x41234125",
"0x51512342",
[1, 2, 3, 4, 5]
],
[ // Event 2: BlockUpdateEvent
"BlockUpdateEvent",
"key_12345",
[
["memory", "tcp://192.168.1.10:6000"],
["memory", "tcp://192.168.1.11:5001"],
["disk", "/data/blocks/block_12345.bin"]
]
],
[ // Event 3: RemoveAllEvent
"RemoveAllEvent"
]
]
]
```
------
### 2.2 Event System Facade Class (kv_event_system.h)
#### 2.2.1 Configuration Item Descriptions
| Category | Parameter Name | Corresponding Field in `KVEventPublisherConfig` Structure | Type | Default Value | Description |
| ---------------------- | ------------------------------------- | --------------------------------------------------------- | ------------- | ----------------- | ------------------------------------------------------------ |
| **Basic Switch** | `enable_kv_event_publish` | None (system switch) | `bool` | `false` | Master switch for event publishing functionality. System starts event publishing only when set to `true`. |
| **Network Config** | `kv_event_publisher_endpoint` | `endpoint` | `std::string` | `"tcp://*:19997"` | ZeroMQ bind address, supports `tcp://`and `ipc://`protocols. |
| | `kv_event_publisher_replay_endpoint` | `replay_endpoint` | `std::string` | `""` | Replay endpoint address. Empty value disables replay functionality. |
| | `kv_event_publisher_topic` | `topic` | `std::string` | `"mooncake"` | Message topic. |
| **Performance Config** | `kv_event_publisher_hwm` | `hwm` | `int` | `100000` | ZeroMQ high-water mark, controls memory buffer size. |
| | `kv_event_publisher_send_interval_ms` | `send_interval` | `uint32_t` | `0` | Send interval (milliseconds), 0 means no delay. |
| | `kv_event_publisher_max_batch_size` | `max_batch_size` | `uint32_t` | `50` | Maximum batch size, affects throughput and latency. |
| **Advanced Config** | `kv_event_publisher_auto_port` | `auto_port` | `bool` | `true` | Automatic port switching, automatically tries other ports if the port is occupied. |
#### 2.2.2 Event Publishing Interface
##### Generic Publishing Method
```c++
// Template method, supports all event types derived from KVCacheEvent
template <DerivedFromKVCacheEvent Event, typename... Args>
std::future<bool> publish(Args&&... args);
```
##### Publishing Different Types of Events
```c++
// 1. Publish block storage event
system.publish<BlockStoreEvent>(
"key_001",
replica_list,
store_info
);
// 2. Publish block update event
system.publish<BlockUpdateEvent>(
"key_001",
updated_replicas
);
// 3. Publish clear all event
system.publish<RemoveAllEvent>();
// 4. Custom event (must inherit from KVCacheEvent)
class CustomEvent : public KVCacheEvent { /* ... */ };
system.publish<CustomEvent>(arg1, arg2);
```
------
## 3. Python API Integration
Since only the `BlockStoreEvent` requires the introduction of field information defined in `StoreEventInfo` passed from the inference engine side, a new default parameter (i.e., the optional parameter `store_event_infos`) has been added only to the interfaces related to `put` operations.
> Currently, only the `batch_put_from_multi_buffers` interface has been adapted accordingly. Other interfaces related to `put` operations will also need adaptation subsequently.
### `StoreEventInfo` Type
**Purpose**
The `StoreEventInfo` struct is used to carry business metadata for KV cache storage events. All fields in this structure are optional; which fields to populate should be determined by the specific business scenario of the upper-layer cache-aware component (such as indexer or router). The Mooncake system acts only as a transparent pipeline for transmitting these fields and will not perform any business logic validation or interpretation of their content.
For detailed API specification, please refer to the Python API Reference documentation: [mooncake-store.md](https://github.com/OpenMooncake/Mooncake/blob/main/docs/source/python-api-reference/mooncake-store.md).
### `batch_put_from_multi_buffers` Interface
The `batch_put_from_multi_buffers` interface has been enhanced to support passing `StoreEventInfo` objects via the optional `store_event_infos` parameter. This allows KV cache storage events to carry additional metadata that can be consumed by downstream components.
For detailed API specification and usage examples, please refer to the Python API Reference documentation: [mooncake-store.md](https://github.com/OpenMooncake/Mooncake/blob/main/docs/source/python-api-reference/mooncake-store.md).

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.
@ -101,30 +100,10 @@ The data structure details of `ReplicateConfig` are as follows:
struct ReplicateConfig {
size_t replica_num{1}; // Total number of replicas for the object
bool with_soft_pin{false}; // Whether to enable soft pin mechanism for this object
bool with_hard_pin{false}; // Whether to enable hard pin (never evicted)
std::string preferred_segment{}; // Preferred segment for allocation
};
```
### Upsert
```C++
tl::expected<void, ErrorCode> Upsert(const ObjectKey& key,
std::vector<Slice>& slices,
const ReplicateConfig& config);
std::vector<tl::expected<void, ErrorCode>> BatchUpsert(
const std::vector<ObjectKey>& keys,
std::vector<std::vector<Slice>>& batched_slices,
const ReplicateConfig& config);
```
`Upsert` inserts `key` if it does not exist and updates the existing object if
it does. It uses the same replication configuration model as `Put`, while
allowing the store to reuse existing placement for in-place updates when the
current layout permits it. `BatchUpsert` performs the same operation for
multiple keys using a shared replication configuration.
### Remove
```C++
@ -536,40 +515,6 @@ The Master Service handles object-related interfaces as follows:
Before writing an object, the Client calls PutStart to request storage space allocation from the Master Service. After completing data writing, the Client calls PutEnd to notify the Master Service to mark the object write as completed.
- Upsert
```C++
tl::expected<std::vector<Replica::Descriptor>, ErrorCode> UpsertStart(
const std::string& key,
const std::vector<size_t>& slice_lengths,
const ReplicateConfig& config);
std::vector<tl::expected<std::vector<Replica::Descriptor>, ErrorCode>>
BatchUpsertStart(const std::vector<std::string>& keys,
const std::vector<std::vector<uint64_t>>& slice_lengths,
const ReplicateConfig& config);
tl::expected<void, ErrorCode> UpsertEnd(
const std::string& key, ReplicaType replica_type);
std::vector<tl::expected<void, ErrorCode>> BatchUpsertEnd(
const std::vector<std::string>& keys);
tl::expected<void, ErrorCode> UpsertRevoke(
const std::string& key, ReplicaType replica_type);
std::vector<tl::expected<void, ErrorCode>> BatchUpsertRevoke(
const std::vector<std::string>& keys);
```
`UpsertStart` / `UpsertEnd` / `UpsertRevoke` mirror the existing put lifecycle
but operate on insert-or-update semantics. If the key does not exist, the flow
behaves like `PutStart`. If the key already exists, the Master may reuse the
current allocation for an in-place update or allocate new space when the object
layout changes. The batch variants provide the same control flow for multiple
keys and are the lower-level primitives used by the high-level `BatchUpsert`
path.
- GetReplicaList
```C++
@ -743,18 +688,6 @@ There are two startup parameters in `master_service` related to the soft pin mec
Notably, soft pinned objects can still be removed using APIs such as `Remove` or `RemoveAll`.
### Hard Pin
For objects that must never be evicted under any circumstances (e.g., model weights, critical metadata), Mooncake Store provides a hard pin mechanism. Unlike soft pin, hard-pinned objects are permanently protected from eviction — they will never be selected as eviction candidates regardless of memory pressure.
Hard pin is set at object creation time through the `with_hard_pin` field in `ReplicateConfig` and cannot be changed afterward. Hard-pinned objects can only be removed explicitly via `Remove` (with force) or `RemoveAll`.
Key differences from soft pin:
- Hard pin never expires. Soft pin status is removed after a configurable TTL if the object is not accessed.
- Hard-pinned objects are completely skipped during eviction. Soft-pinned objects may still be evicted when no other candidates are available.
- Hard pin is immutable once set. Soft pin status is automatically refreshed on access.
### Zombie Object Cleanup
If a Client crashes or experiences a network failure after sending a `PutStart` request but before it can send the corresponding `PutEnd` or `PutRevoke` request to the Master, the object initiated by `PutStart` enters a "zombie" state—rendering it neither usable nor deletable. The existence of such "zombie objects" not only consumes storage space but also prevents subsequent `Put` operations on the same keys. To mitigate these issues, the Master records the start time of each `PutStart` request and employs two timeout thresholds—`put_start_discard_timeout` and `put_start_release_timeout`—to clean up zombie objects.
@ -779,7 +712,6 @@ The preferred segment allocation feature is implemented through the `AllocationS
struct ReplicateConfig {
size_t replica_num{1}; // Total number of replicas for the object
bool with_soft_pin{false}; // Whether to enable soft pin mechanism for this object
bool with_hard_pin{false}; // Whether to enable hard pin (never evicted)
std::string preferred_segment{}; // Preferred segment for allocation
};
```
@ -1043,13 +975,3 @@ When to bump the version:
* **Major version (X.0.0)**: For breaking API changes, major architectural changes, or significant new features that affect backward compatibility
* **Minor version (0.X.0)**: For new features, API additions, or notable improvements that maintain backward compatibility
* **Patch version (0.0.X)**: For bug fixes, performance optimizations, or minor improvements that don't affect the API
---
:::{toctree}
:caption: Related Design Docs
:maxdepth: 1
ssd-offload
:::

View File

@ -1,245 +0,0 @@
# SSD Offload Design
## Overview
Mooncake Store supports offloading KV cache objects from distributed memory to local SSD. This extends the effective cache capacity beyond DRAM limits at lower cost, while preserving the performance characteristics of the hot path through zero-copy RDMA-based memory transfers.
SSD offload is implemented as a background subsystem within the **real client** process. It is transparent to the application: a `Put` that would otherwise be evicted from memory is persisted to disk, and a `Get` that finds no memory replica automatically falls back to reading from SSD.
---
## Architecture
```
┌─────────────────────────────────────────────────────────┐
│ Application (vLLM, etc.) │
└──────────────────────────┬──────────────────────────────┘
│ MooncakeDistributedStore API
┌─────────────────────────────────────────────────────────┐
│ Real Client │
│ │
│ ┌──────────────────────────────────────────────────┐ │
│ │ FileStorage │ │
│ │ ┌────────────┐ ┌──────────────────────────┐ │ │
│ │ │ Heartbeat │ │ ClientBuffer (staging) │ │ │
│ │ │ Thread │ └──────────────────────────┘ │ │
│ │ └─────┬──────┘ │ │
│ │ │ offload / load │ │
│ │ ▼ │ │
│ │ ┌─────────────────────────────────────────┐ │ │
│ │ │ StorageBackendInterface │ │ │
│ │ │ ┌───────────┐ ┌──────────┐ ┌────────┐ │ │ │
│ │ │ │ Bucket │ │FilePerKey│ │Offset │ │ │ │
│ │ │ │ Backend │ │ Backend │ │Alloc. │ │ │ │
│ │ │ └───────────┘ └──────────┘ └────────┘ │ │ │
│ │ └─────────────────────────────────────────┘ │ │
│ └──────────────────────────────────────────────────┘ │
│ │
│ ┌──────────────────────────────────────────────────┐ │
│ │ In-memory distributed KV cache │ │
│ │ (Transfer Engine / RDMA) │ │
│ └──────────────────────────────────────────────────┘ │
└─────────────────────────────────────────────────────────┘
Local SSD / NVMe
```
The key components are:
- **FileStorage**: The top-level coordinator. It owns the storage backend, a staging buffer (`ClientBuffer`), and background threads for heartbeating and buffer garbage collection.
- **StorageBackendInterface**: An abstract interface implemented by three backends (see below). Responsible for the actual on-disk layout and I/O.
- **Heartbeat thread**: Periodically contacts the master. The master returns a list of objects to offload; the heartbeat thread writes them to SSD and notifies the master of completion.
- **ClientBuffer**: A pre-registered, O_DIRECT-aligned staging area used for zero-copy reads from SSD back into application memory.
---
## Data Flow
### Offload (memory → SSD)
The offload path is driven entirely by the heartbeat thread inside `FileStorage`. No write path from the application is involved.
```
Heartbeat Thread Master Local Memory Segment
│ │ │
│─OffloadObjectHB ───▶│ │
│◀─ {key→size} map ───│ (objects to evict from │
│ │ memory to SSD) │
│ │ │
│─ BatchQuery(keys) ───────────────────────────────▶│
│◀─ {key→Slice} ────────────────────────────────────│
│ │ │
│ [PrepareEviction: remove old buckets, notify master via BatchEvictDiskReplica]
│─ BatchEvictDiskReplica(evicted_keys) ────────────▶│ (master removes stale replicas)
│ [FinalizeEviction: delete evicted files]
│ │ │
│ BatchOffload(slices) → StorageBackend → SSD │
│ │ │
│─ NotifyOffloadSuccess(keys, metadata) ───────────▶│
│ │ (master adds LOCAL_DISK │
│ │ replica to object entry) │
```
Step by step:
1. **Heartbeat**: The heartbeat thread wakes up every `MOONCAKE_OFFLOAD_HEARTBEAT_INTERVAL_SECONDS` seconds and calls `client_->OffloadObjectHeartbeat(enable_offloading_, offloading_objects)`. The master replies with a map of `{key → size}` for objects it has selected to evict from memory.
2. **Read slices from memory**: `FileStorage::OffloadObjects` groups the keys into buckets (for `BucketStorageBackend`) and calls `BatchQuerySegmentSlices` to obtain `{key → Slice}` from the local memory segment via `client_->BatchQuery`.
3. **Eviction** (if capacity limit is set): Before writing, `PrepareEviction` removes old buckets from metadata under the exclusive lock and collects their keys. The `eviction_handler` callback calls `client_->BatchEvictDiskReplica` to notify the master in a single RPC. `FinalizeEviction` then deletes the corresponding files.
4. **Write to SSD**: `StorageBackend::BatchOffload` serializes and writes the key-value data to disk.
5. **Notify master**: On success, the `complete_handler` calls `client_->NotifyOffloadSuccess(keys, metadatas)`. The master adds a `LOCAL_DISK` replica entry (carrying the real client's RPC address as `transport_endpoint`) to the object's replica list.
### Load (SSD → memory)
The load path involves three parties: the **requesting client**, the **target client** that holds the SSD data, and the **Transfer Engine** for zero-copy data movement.
```
Requesting Client Target Client Master
│ │ │
│─ BatchGet(keys) ──────────────────────────────────────────▶│
│◀─ QueryResult {replicas: [LOCAL_DISK(rpc_addr)]} ──────────│
│ │ │
│ (no memory replica available) │ │
│─ batch_get_offload_object(keys, sizes) ───────────────────▶│
│ │ │
│ FileStorage::BatchGet │
│ → StorageBackend::BatchLoad │
│ → read from SSD into ClientBuffer │
│ │ │
│◀─ BatchGetOffloadObjectResponse ───│ │
│ {batch_id, pointers[], transfer_engine_addr, gc_ttl_ms} │
│ │ │
│─ Transfer Engine: BatchGetOffloadObject ──────────────────▶│
│ (RDMA/TCP: pull data from ClientBuffer into app memory) │
│◀─ done ────────────────────────────│ │
│ │ │
│─ release_offload_buffer(batch_id) ────────────────────────▶│
│ │ (free ClientBuffer slot)│
```
Step by step:
1. **Query master**: The requesting client calls `client_->BatchGet(keys, ...)` to query the master for replica locations. If the object has been offloaded, the master returns a `LOCAL_DISK` replica descriptor containing the target client's RPC address (`transport_endpoint`).
2. **RPC to target client**: The requesting client calls `batch_get_offload_object(keys, sizes)` on the target client identified by `transport_endpoint`. The target client calls `FileStorage::BatchGet`, which allocates slots in `ClientBuffer` and reads the requested objects from SSD via `StorageBackend::BatchLoad`.
3. **Response with buffer pointers**: The target client returns a `BatchGetOffloadObjectResponse` containing `batch_id`, a list of buffer `pointers` (addresses within `ClientBuffer`), the Transfer Engine address, and `gc_ttl_ms` (the buffer lease TTL).
4. **Zero-copy transfer**: The requesting client invokes `client_->BatchGetOffloadObject(transfer_engine_addr, keys, pointers, slices)`, which uses the Transfer Engine (RDMA or TCP) to pull the data directly from the target client's `ClientBuffer` into the application's target memory (DRAM or VRAM). No intermediate copy is made on the requesting client side.
5. **Release buffer**: After the transfer completes, the requesting client immediately calls `release_offload_buffer(batch_id)` on the target client to free the `ClientBuffer` slots. If the transfer takes longer than `gc_ttl_ms`, the buffer GC thread reclaims the slot automatically as a fallback.
---
## Storage Backends
### BucketStorageBackend (default)
Objects are grouped into **buckets** before being written to disk. Each bucket produces two files:
- **`.bucket`** — binary data file containing serialized key-value records
- **`.meta`** — metadata file describing the keys and byte offsets within the data file
Bucket IDs are monotonically increasing timestamps with a sequence suffix, so `buckets_` (a `std::map<int64_t, BucketMetadata>`) is always ordered by creation time.
**Grouping strategy** (`GroupOffloadingKeysByBucket`): objects are accumulated into a bucket until either `bucket_size_limit` (default 256 MB) or `bucket_keys_limit` (default 500) is reached. Objects that do not fill a complete bucket are held in `ungrouped_offloading_objects_` and retried on the next heartbeat.
**In-flight read tracking**: A `BucketReadGuard` RAII object increments `BucketMetadata::inflight_reads_` on construction and decrements it on destruction. This allows safe deletion of bucket files even when concurrent reads are in progress.
### StorageBackendAdaptor (FilePerKey)
Each object is stored as an individual file. The file path is derived from the key via a two-level hash-sharded directory structure to avoid large flat directories. This backend is simple and easy to inspect but does not scale well to millions of objects.
### OffsetAllocatorStorageBackend
A single pre-allocated file (`kv_cache.data`) is shared by all objects. Space within the file is managed by an `OffsetAllocator`. Metadata is sharded across 1024 independent maps to reduce lock contention under high concurrency. Records follow the layout `[key_len: u32 | value_len: u32 | key | value]`.
---
## 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`).
### Policies
| Policy | Candidate selection |
|--------|---------------------|
| `FIFO` | `buckets_.begin()` — always the oldest bucket, since `buckets_` is ordered by bucket ID |
| `LRU` | `std::min_element` over `BucketMetadata::last_access_ns_` — the bucket with the smallest last-read timestamp |
`last_access_ns_` is an atomic `int64_t` updated on every `BatchLoad` with relaxed ordering. Buckets that have never been read have `last_access_ns_ == 0` and are therefore always evicted first under LRU, giving FIFO-among-unread semantics.
### Two-phase eviction protocol
Eviction is split into two phases to ensure that the master is notified before files are deleted, and that no in-flight reads are interrupted.
**Phase 1 — `PrepareEviction(required_size)`** (called under exclusive lock):
1. Repeatedly call `SelectEvictionCandidate()` until `total_size_ + required_size <= max_total_size`.
2. For each selected bucket: remove it from `buckets_` and `object_bucket_map_`, subtract its size from `total_size_`.
3. Collect all evicted keys and bucket metadata into a `PendingEviction` struct and return it — no file I/O at this point.
**Between phases** — notify master:
The caller invokes the `eviction_handler` callback with the full list of evicted keys. The handler calls `MasterClient::BatchEvictDiskReplica`, which sends a single RPC to the master to remove the disk replicas for all evicted keys atomically.
**Phase 2 — `FinalizeEviction(pending)`** (called after master notification):
For each evicted bucket:
1. Spin-wait (with a 10-second timeout) until `inflight_reads_ == 0`.
2. Evict any stale file-handle cache entries.
3. Delete the `.bucket` and `.meta` files.
This ordering guarantees:
- The master never serves a stale disk-replica location for a file that has already been deleted.
- Ongoing reads complete successfully before their files are removed.
- Freed disk space is available for the incoming write before `WriteBucket` is called.
---
## 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.
### Thread-local rings (`SharedUringRing`)
Each thread owns exactly one `io_uring` ring, stored in `thread_local` storage. This means:
- **No mutex between threads.** Each ring is accessed only by its owning thread, so concurrent I/O from multiple threads is fully parallel with zero synchronization overhead.
- **Within-thread batching.** Multiple SQEs can be enqueued before calling `io_uring_submit_and_wait`, exposing NVMe queue depth > 1 within a single thread. `batch_read` exploits this to issue up to `QUEUE_DEPTH` (32) independent reads in one submission.
- **File-descriptor registration is omitted.** The per-I/O `fdget()` overhead (~50 ns) is negligible compared to the lock contention (> 1 ms) the old global ring imposed, so `IOSQE_FIXED_FILE` is not used.
Rings are initialized lazily on first use and destroyed when the thread exits. If ring initialization fails (e.g., kernel too old), the backend falls back gracefully to POSIX I/O.
### Fixed-buffer registration
The `ClientBuffer` (the staging buffer used for SSD reads) is registered with io_uring as a **fixed buffer** via `io_uring_register_buffers`. When a read destination falls within the registered region, the backend uses `io_uring_prep_read_fixed` instead of `io_uring_prep_read`, which avoids a per-I/O `mmap`/`munmap` in the kernel and reduces system-call overhead.
Buffer registration is global but applied **lazily per thread**: `g_buf` stores the base address and length atomically; each thread-local ring calls `ensure_buf_registered()` on its first I/O and registers the buffer independently. This avoids a global barrier at startup.
To prevent `io_uring`'s `FOLL_LONGTERM` page pinning from failing on systems with Transparent Huge Pages (THP) enabled, `MADV_NOHUGEPAGE` is applied to the buffer region before registration. This forces the kernel to back the range with 4 KB pages, making long-term pinning reliable regardless of system THP policy.
### O_DIRECT and alignment
`UringFile` supports an optional `O_DIRECT` mode. When enabled:
- All file descriptors are opened with `O_DIRECT`.
- Buffers, lengths, and offsets must be aligned to 4 KB (`ALIGNMENT_ = 4096`).
- For unaligned writes (e.g., metadata serialized into a `std::string`), the backend allocates a temporary aligned bounce buffer via `posix_memalign`, copies the data, performs the aligned write, and frees the bounce buffer.
- `read_aligned` and `write_aligned` are the primary I/O paths; they assert alignment constraints and delegate directly to the ring.
### I/O operations
| Method | Description |
|--------|-------------|
| `read` / `write` | Contiguous read or write, chunked into up to `QUEUE_DEPTH` SQEs per submission |
| `read_aligned` / `write_aligned` | Same as above but with alignment preconditions for O_DIRECT |
| `batch_read` | Submits multiple independent reads (different offsets) in batches of up to `QUEUE_DEPTH`, maximizing NVMe queue utilization |
| `vector_read` / `vector_write` | Scatter/gather I/O: one SQE per `iovec`, submitted in batches |
| `datasync` | Issues `IORING_FSYNC_DATASYNC` and waits for completion |
### Integration with storage backends
- **BucketStorageBackend**: uses `UringFile` for both bucket data files and metadata files when `use_uring_` is set. A file-handle cache (`file_cache_`) avoids repeated `open`/`close` for hot buckets. On eviction, the cache entry is explicitly removed before the file is deleted to prevent stale handles.
- **OffsetAllocatorStorageBackend**: opens the single pre-allocated data file with `O_DIRECT` and `UringFile`, and uses `GetFileInstance()` to expose the file handle for external buffer registration.
- **StorageBackendAdaptor** (FilePerKey): uses `UringFile` for reads when `use_uring_` is set; writes use POSIX paths.
## Metadata Recovery on Restart
On startup, `FileStorage::Init` calls `StorageBackend::ScanMeta`, which reads all on-disk metadata and invokes a callback for each discovered object. The callback calls `MasterClient::NotifyOffloadSuccess` to re-register the objects with the master. This restores the full disk-replica view without any application-level intervention.

View File

@ -2,8 +2,6 @@
The source code path for Ascend Transport is `Mooncake/mooncake-transfer-engine/src/transport/ascend_transport`, which also includes automated build scripts and the README file.
**ASCEND TRANSPORT is scheduled for deprecation, please use [ASCEND DIRECT TRANSPORT](./ascend_direct_transport.md) on ASCEND platform. **
## Overview
Ascend Transport is a high-performance zero-copy NPU data transfer library with one-sided semantics, directly compatible with Mooncake Transfer Engine. To compile and use the Ascend Transport library, please set the `USE_ASCEND` flag to `"ON"` in the `mooncake-common/common.cmake` file.
@ -144,7 +142,7 @@ Therefore, in testing:
Watch the log produced by `mooncake-transfer-engine/src/transfer_engine.cpp`; you should see a line similar to
```
Transfer Engine RPC using <protocol> listening on <IP>:<actual-port>
```
```
Note the **actual port** the target node is listening on.
2. **Edit the initiators launch command**:

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,46 @@ 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 \
libyaml-cpp-dev \
libgtest-dev \
pybind11-dev \
python3-dev
# Install yalantinglibs (required)
cd /tmp
git clone https://github.com/alibaba/yalantinglibs.git
cd yalantinglibs
mkdir build && cd build
cmake .. -DCMAKE_INSTALL_PREFIX=/usr/local
make -j$(nproc)
sudo make install
```
## Building Mooncake with EFA Support
### 1. Clone the Repository
```bash
git clone https://github.com/kvcache-ai/Mooncake.git
cd Mooncake
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 +76,12 @@ 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 +102,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 +184,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 +193,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 +287,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 +303,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 +351,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

@ -18,7 +18,6 @@ pip install mooncake-transfer-engine-non-cuda
📦 **Package Details**: [https://pypi.org/project/mooncake-transfer-engine-non-cuda/](https://pypi.org/project/mooncake-transfer-engine-non-cuda/)
> **Note**: The CUDA version includes Mooncake-EP and GPU topology detection, requiring CUDA 12.1+. The non-CUDA version is for environments without CUDA dependencies.
> **Note**: MLU support is currently source-build only. If you need Cambricon MLU memory support, install Neuware and build with `-DUSE_MLU=ON`.
## Automatic
@ -113,43 +112,8 @@ pip install mooncake-transfer-engine-non-cuda
```bash
export LIBRARY_PATH=$LIBRARY_PATH:/usr/local/musa/lib
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/musa/lib
```
4. If you want to compile Cambricon MLU support, first install the Cambricon Neuware SDK. After that:
1) Export `NEUWARE_HOME` or pass `-DNEUWARE_ROOT=/path/to/neuware` to CMake
2) Configure `LIBRARY_PATH` and `LD_LIBRARY_PATH` to ensure linking of `cnrt`, `cndrv`, and other Neuware libraries during compilation:
```bash
export NEUWARE_HOME=/usr/local/neuware
export LIBRARY_PATH=$LIBRARY_PATH:${NEUWARE_HOME}/lib64
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${NEUWARE_HOME}/lib64
```
If your Neuware installation lives outside the default include/library layout, you can also pass:
```bash
cmake .. -DUSE_MLU=ON \
-DMLU_INCLUDE_DIR=/path/to/neuware/include \
-DMLU_LIB_DIR=/path/to/neuware/lib64
```
For Cambricon MLU builds, enable the MLU backend explicitly:
```bash
cmake .. -DUSE_MLU=ON -DNEUWARE_ROOT=${NEUWARE_HOME:-/usr/local/neuware}
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:
```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)
6. Install yalantinglibs
4. Install yalantinglibs
```bash
git clone https://github.com/alibaba/yalantinglibs.git
cd yalantinglibs
@ -159,7 +123,7 @@ pip install mooncake-transfer-engine-non-cuda
make install
```
7. In the root directory of this project, run the following commands:
5. In the root directory of this project, run the following commands:
```bash
mkdir build
cd build
@ -167,7 +131,7 @@ pip install mooncake-transfer-engine-non-cuda
make -j
```
8. Install Mooncake python package and mooncake_master executable
6. Install Mooncake python package and mooncake_master executable
```bash
make install
```
@ -187,25 +151,15 @@ cd /Mooncake-main/build/mooncake-transfer-engine/example
## Advanced Compile Options
The following options can be used during `cmake ..` to specify whether to compile certain components of Mooncake.
- `-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_MNNVL=[ON|OFF]`: Enable Multi-Node NVLink transport support, default is OFF. **Note:** `-DUSE_CUDA` is required when `-DUSE_MNNVL` is 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_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`.
- `-DMLU_INCLUDE_DIR=/path/to/include`: Override the Neuware include directory when `-DUSE_MLU=ON`.
- `-DMLU_LIB_DIR=/path/to/lib64`: Override the Neuware library directory when `-DUSE_MLU=ON`.
- `-DUSE_EFA=[ON|OFF]`: Enable AWS Elastic Fabric Adapter transport via libfabric. **Default: OFF.** See [EFA Transport](../design/transfer-engine/efa_transport.md) for details.
- `-DUSE_INTRA_NVLINK=[ON|OFF]`: Enable intranode nvlink transport
- `-DUSE_CXL=[ON|OFF]`: Enable CXL support
- `-DWITH_STORE=[ON|OFF]`: Build Mooncake Store component
- `-DWITH_P2P_STORE=[ON|OFF]`: Enable Golang support and build P2P Store component, require go 1.23+
- `-DWITH_WITH_RUST_EXAMPLE=[ON|OFF]`: Enable Rust support
- `-DWITH_EP=[ON|OFF]`: Build the EP (Expert Parallelism) and PG Python extensions for CUDA. Requires CUDA toolkit and PyTorch. Use `-DEP_TORCH_VERSIONS="2.9.1"` (semicolon-separated) to build for specific PyTorch versions, or leave empty to use the currently-installed torch. The CUDA version is detected automatically. **Default: OFF.**
- `-DUSE_REDIS=[ON|OFF]`: Enable Redis-based metadata service
- `-DUSE_HTTP=[ON|OFF]`: Enable Http-based metadata service
- `-DUSE_ETCD=[ON|OFF]`: Enable etcd-based metadata service, require go 1.23+

View File

@ -200,8 +200,6 @@ mooncake_master \
```
This exposes the metadata endpoint at `http://<host>:<port>/metadata`.
If the master runs in a container and its IP is dynamic, set `--rpc_interface=<ifname>` such as `--rpc_interface=eth0`. Mooncake Master will resolve the current IPv4 address from that interface at startup instead of relying on a fixed `--rpc_address`.
Optional: Use the free-ratio-first allocation strategy for better load balancing across segments with different sizes or utilization:
```bash
@ -245,4 +243,4 @@ store.close()
### More Examples and Documentation
Please refer to the [Mooncake Store Python API](../python-api-reference/mooncake-store.md), [Mooncake Store](../design/mooncake-store.md) and [Mooncake Store Deployment & Operations Guide](../deployment/mooncake-store-deployment-guide.md) for more examples and documentation.
Please refer to the [Mooncake Store Python API](../python-api-reference/mooncake-store.md), [Mooncake Store](../design/mooncake-store.md) and [Mooncake Store Deployment & Operations Guide](../deployment/mooncake-store-deployment-guide.md) for more examples and documentation.

View File

@ -56,7 +56,7 @@ export MOONCAKE_PROTOCOL="tcp"
### RDMA (Recommended for Production)
**Description:** Remote Direct Memory Access protocol providing high-performance, low-latency data transfer with minimal CPU overhead. Supports accelerator-aware memory registration, including NVIDIA GPUDirect RDMA for CUDA buffers and Cambricon MLU buffers when built with Neuware.
**Description:** Remote Direct Memory Access protocol providing high-performance, low-latency data transfer with minimal CPU overhead. Supports GPUDirect RDMA for zero-copy GPU memory transfers.
**Hardware Support:**
- InfiniBand
@ -64,7 +64,6 @@ export MOONCAKE_PROTOCOL="tcp"
- eRDMA (Elastic RDMA)
- NVIDIA GPUDirect RDMA
- Non-NVIDAI GPUDirect RDMA (e.g., Intel E810 RDMA NIC)
- Cambricon MLU memory via Neuware (`-DUSE_MLU=ON`)
**Use When:**
- High-performance networking is required
@ -73,8 +72,6 @@ export MOONCAKE_PROTOCOL="tcp"
**Note:** If no RDMA HCA (Host Channel Adapter) is detected on the system, the Transfer Engine will automatically fall back to TCP protocol for compatibility.
**MLU Note:** Cambricon MLU support uses the standard `rdma` data path. There is no separate `mlu` protocol string. To enable MLU memory detection, topology discovery, and DMA-BUF based registration, build Transfer Engine with `-DUSE_MLU=ON` and make Neuware available through `NEUWARE_HOME` or `NEUWARE_ROOT`.
**Configuration:**
```python
# Python API - With specific device
@ -324,7 +321,6 @@ export MOONCAKE_LOCAL_HOSTNAME="node1"
| Cloud Environments | tcp or rdma (if available) | Check cloud provider support |
| Multi-tier Storage | rdma + nvmeof | Combine protocols for different layers |
| AMD GPU Clusters | rdma + hip | Use HIP for local GPU communication |
| Cambricon MLU Clusters | rdma | Build with `-DUSE_MLU=ON`; MLU uses the normal RDMA protocol |
| Ascend NPU Clusters | rdma + ascend | Use Ascend for NPU-specific operations |
## Troubleshooting

Binary file not shown.

Before

Width:  |  Height:  |  Size: 73 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 140 KiB

View File

@ -27,7 +27,6 @@ This repository also hosts its technical report and the open-sourced traces.
<h2 id="updates">🔄 Updates</h2>
- **Mar 19, 2026**: [TorchSpec: Speculative Decoding Training at Scale](https://pytorch.org/blog/torchspec-speculative-decoding-training-at-scale) is [open sourced](https://github.com/torchspec-project/TorchSpec), using Mooncake to decouple inference and training via efficient hidden states management.
- **Feb 12, 2026**: [Mooncake Joins PyTorch Ecosystem](https://pytorch.org/blog/mooncake-joins-pytorch-ecosystem/) We are thrilled to announce that Mooncake has officially joined the PyTorch Ecosystem!
- **Jan 28, 2026**: [FlexKV](https://github.com/taco-project/FlexKV), a distributed KV store and cache system from Tencent and NVIDIA in collaboration with the community, now supports [distributed KVCache reuse](https://github.com/taco-project/FlexKV/blob/main/docs/dist_reuse/README_en.md) with the Mooncake Transfer Engine.
- **Dec 23, 2025**: SGLang introduces [Encode-Prefill-Decode (EPD) Disaggregation](https://lmsys.org/blog/2026-01-12-epd/) with Mooncake as a transfer backend. This integration allows decoupling compute-intensive multimodal encoders (e.g., Vision Transformers) from language model nodes, utilizing Mooncake's RDMA engine for zero-copy transfer of large multimodal embeddings.
@ -86,7 +85,6 @@ performance/vllm-benchmark-results-v1
performance/sglang-hicache-benchmark-results-v1
performance/vllm-v1-support-benchmark
performance/allocator-benchmark-result
performance/ssd-offload-benchmark-results
:::
% API Documentation
@ -130,7 +128,7 @@ troubleshooting/troubleshooting
:::{toctree}
:caption: Deployment
:maxdepth: 2
:maxdepth: 1
deployment/mooncake-store-deployment-guide
:::

View File

@ -1,157 +0,0 @@
# Mooncake SSD Offload Benchmark
This benchmark measures the performance benefit of Mooncake's SSD offload feature in multi-turn conversation scenarios. In the test, multiple clients send requests concurrently, each simulating a multi-round dialogue where every new round appends the previous context.
We compare four storage configurations for the KV cache:
* **GPU only**: KV cache resides entirely in GPU memory.
* **(HiCache L1) + L2**: KV cache spans GPU and host memory via HiCache's two-level hierarchy.
* **(HiCache L1 + L2) + Mooncake**: KV cache is further extended into an 80GB Mooncake distributed memory pool.
* **(HiCache L1 + L2) + Mooncake + SSD**: On top of the above, SSD offload is enabled so that evicted cache entries are written to local NVMe storage rather than discarded.
The benchmark targets the prefill stage and reports two primary metrics: Time-To-First-Token (TTFT) and input token throughput.
## Benchmark Result
![overall performance](../image/ssd_offload_overall.png)
The figure above summarizes the end-to-end results on a single DGX node (8 × A100-SXM4-40GB, dual RDMA NICs). Enabling SSD offload cuts average TTFT by **57%** relative to GPU only and by **34%** relative to Mooncake without SSD, while delivering a **2.4×** improvement in input token throughput.
![per-turn performance](../image/ssd_offload_per_turn.png)
To better understand where the gains come from, we break down TTFT and cache hit rate by conversation round. The output length is fixed to 1 token so that decode overhead does not obscure prefill differences.
During the first six rounds the 80GB memory pool is large enough, so `+ Mooncake` and `+ Mooncake + SSD` behave identically — both sustain hit rates above 80%.
The divergence appears in round 7. Once the accumulated KV cache exceeds memory capacity, `+ Mooncake` must evict entries and its hit rate plunges from 83% to 36%, pushing TTFT from 6s to 16s. With SSD offload, those evicted entries survive on disk and remain retrievable; the hit rate stays above 84% through round 8, and TTFT remains at 9.4s — roughly half the latency of Mooncake without SSD.
Note that a slight increase in TTFT is visible in round 8 with SSD offload (9.4s vs 7.4s in round 7), reflecting the additional latency of reading evicted entries from NVMe storage rather than RDMA memory. This overhead is modest compared to the alternative of re-computing evicted KV cache from scratch.
This demonstrates that SSD offload turns local NVMe drives into a cost-effective extension of the cache hierarchy. In production, where long conversations and high concurrency are common, this prevents the sharp performance cliff that occurs when DRAM-based caching alone is exhausted.
## Benchmark Setup
### DGX Server
**Experimental Environment**
- GPU: 8 × NVIDIA A100-SXM4-40GB
- Network: Dual RDMA NICs (ibp12s0, ibp75s0), InfiniBand 4X HDR 200 Gb/s each
- Storage: 5 × Samsung NVMe SSDs in RAID0 — 3 × PM1733 3.84TB (PCIe Gen4, 7,000 MB/s seq read each) + 2 × PM983 1.92TB (PCIe Gen3, 3,000 MB/s seq read each). Aggregate theoretical sequential read bandwidth: ~27 GB/s. Mounted at /mnt/data (~14TB usable), used as the SSD offload target.
- Model: Qwen3-8B
**Benchmark Script:**
We used SGLang's [multiturn benchmark](https://github.com/sgl-project/sglang/blob/main/benchmark/hicache/bench_multiturn.py) for the evaluation.
```bash
python3 benchmark/hicache/bench_multiturn.py \
--model-path $MODEL_PATH \
--host 127.0.0.1 \
--port 8189 \
--disable-random-sample \
--output-length 1 \
--request-length 4096 \
--num-clients 20 \
--num-rounds 10 \
--max-parallel 4 \
--request-rate 16 \
--ready-queue-policy random \
--disable-auto-run \
--enable-round-barrier
```
**GPU Only:**
```bash
python3 -m sglang.launch_server \
--model-path $MODEL_PATH \
--tp 1 \
--page-size 64 \
--attention-backend triton
```
**HiCache L1 + L2:**
```bash
python3 -m sglang.launch_server \
--model-path $MODEL_PATH \
--tp 1 \
--page-size 64 \
--attention-backend triton \
--enable-hierarchical-cache \
--hicache-ratio 2
```
**L1 + L2 + Mooncake:**
Mooncake master and client must be started before launching the SGLang server.
```bash
# Start Mooncake master
mooncake_master \
-http_metadata_server_port=8081 \
-metrics_port=9004 \
-logtostderr
# Start Mooncake client (requires root)
# Total Distributed Memory Pool: 80GB
mooncake_client \
--host=127.0.0.1 \
--global_segment_size=80GB \
--master_server_address=localhost:50051 \
--metadata_server=P2PHANDSHAKE \
--protocol=rdma \
--device_names=ibp12s0,ibp75s0 \
--port=50052 \
--logtostderr
```
```bash
MOONCAKE_MASTER="127.0.0.1:50051" \
MOONCAKE_GLOBAL_SEGMENT_SIZE=0 \
MOONCAKE_PROTOCOL="rdma" \
MOONCAKE_DEVICE="ibp12s0,ibp75s0" \
python3 -m sglang.launch_server \
--model-path $MODEL_PATH \
--tp 1 \
--page-size 64 \
--attention-backend triton \
--enable-hierarchical-cache \
--hicache-ratio 2 \
--hicache-storage-prefetch-policy wait_complete \
--hicache-mem-layout page_first_direct \
--hicache-storage-backend mooncake
```
**L1 + L2 + Mooncake + SSD:**
Compared to the previous configuration, the only change is enabling SSD offload on both master and client. A 20GB local buffer absorbs write bursts before flushing to SSD.
```bash
# Start Mooncake master with offload enabled
mooncake_master \
-enable_offload=true \
-http_metadata_server_port=8081 \
-metrics_port=9004 \
-logtostderr
# Start Mooncake client with offload enabled (requires root)
# Total Distributed Memory Pool: 80GB
# 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_client \
--host=127.0.0.1 \
--global_segment_size=80GB \
--master_server_address=localhost:50051 \
--metadata_server=P2PHANDSHAKE \
--protocol=rdma \
--device_names=ibp12s0,ibp75s0 \
--enable_offload=true \
--port=50052 \
--logtostderr
```
The SGLang server launch command is identical to `L1 + L2 + Mooncake`.

View File

@ -1,163 +0,0 @@
# Mooncake KVCache Storage Benchmark
High-performance KVCache storage benchmark tool based on Mooncake Store architecture.
## Overview
Evaluates I/O performance of KVCache storage systems using:
- Single large file (100GB) with offset-based block management
- Prefix caching simulation with hash-based block lookup
- Timestamp-based request replay for realistic testing
- Comprehensive metrics: latency, bandwidth, hit rates
## Test Flow
1. **Load Traces**: Read request sequences from JSONL files (`FAST25-release/traces`)
2. **Process Requests**: For each request, check hash_id prefix cache hits/misses
3. **Perform I/O**: Read cached blocks from disk, write new blocks to storage
4. **Collect Metrics**: Track latency, bandwidth, and cache hit rates
## Quick Start
```bash
# Quick test (100 requests, no timestamp replay)
python storage_benchmark.py --scenario=toolagent --max-requests=100
# Test with large model preset (Llama-3.1-405B)
python storage_benchmark.py --scenario=toolagent --model=llama-3.1-405b --max-requests=100
# Test with Deepseek V3 (extra large model)
python storage_benchmark.py --scenario=toolagent --model=deepseek-v3 --max-requests=100
# Realistic replay (with timestamps, 10x speed)
python storage_benchmark.py --scenario=toolagent --max-requests=1000 \
--replay-timestamps --time-scale=0.1
# Test all scenarios with replay
python storage_benchmark.py --scenario=all --time-scale=1.0
```
## Command-Line Options
| Option | Description | Default |
|--------|-------------|---------|
| `--trace-dir` | Trace files directory | `../FAST25-release/traces` |
| `--scenario` | Test scenario: `conversation`, `synthetic`, `toolagent`, `all` | `toolagent` |
| `--storage-dir` | Storage directory | `/tmp/mooncake_bench` |
| `--model` | Model preset (overrides `--bytes-per-token`) | `default` |
| `--bytes-per-token` | Bytes per token (2048 for 7B FP16) | `2048` |
| `--max-requests` | Maximum requests per scenario (unlimited if not specified) | `None` |
| `--max-blocks` | Maximum number of blocks | `100000` |
| `--replay-timestamps` | Enable timestamp replay | `False` |
| `--time-scale` | Time scaling factor (1.0 = real-time, 0.1 = 10x faster) | `1.0` |
## Model Presets
The tool includes presets for popular LLM models with accurate KVCache sizes based on the [LMCache KVCache Calculator](https://lmcache.ai/kv_cache_calculator.html).
| Model | Bytes/Token | Size | Notes |
|-------|-------------|------|-------|
| **Small Models (7B-13B)** |
| `llama-3-8b` | 128 | 128 B/token | GQA optimized |
| `mistral-7b` | 128 | 128 B/token | GQA optimized |
| `qwen-14b` | 40 | 40 B/token | GQA optimized |
| `gemma-7b` | 224 | 224 B/token | |
| `llama-2-7b` | 512 | 512 B/token | |
| `llama-2-13b` | 800 | 800 B/token | |
| **Large Models (70B-405B)** |
| `llama-2-70b` | 320 | 320 B/token | GQA optimized |
| `llama-3-70b` | 320 | 320 B/token | GQA optimized |
| `mixtral-8x7b` | 128 | 128 B/token | GQA optimized |
| `mixtral-8x22b` | 224 | 224 B/token | GQA optimized |
| `qwen-72b` | 320 | 320 B/token | GQA optimized |
| `qwen-110b` | 320 | 320 B/token | GQA optimized |
| `llama-3.1-405b` | 516018 | ~504 KB/token | Very large KVCache |
| **Extra Large Models** |
| `glm-4.6` | 156991 | ~153 KB/token | |
| `deepseek-v3` | 1749384 | ~1.67 MB/token | Largest KVCache |
| **Default** |
| `default` | 2048 | 2 KB/token | Legacy 7B FP16 |
**Usage**: `--model=llama-3.1-405b` (overrides `--bytes-per-token`)
## Test Scenarios
- **`conversation`**: Write-intensive workload (dialogue patterns)
- **`synthetic`**: Read-intensive workload (cached patterns)
- **`toolagent`**: Balanced read/write mix (tool use patterns)
## Output Example
```
================================================================================
Mooncake KVCache Storage Benchmark
================================================================================
Using model preset: llama-3.1-405b (516018 bytes/token, ~504.0 KB/token)
[1/1] toolagent_trace.jsonl
================================================================================
[Performance Overview]
Total Requests: 100
Queries Per Second (QPS): 14.45
Cache Hit Rate: 24.27%
Write Ratio: 75.73%
Total Blocks: 1,949
Read Blocks: 473
Write Blocks: 1,476
Prefix Hits: 376
[Latency Analysis]
Request Latency (End-to-End): Avg=69.18ms, P50=15.49ms, P95=239.99ms, P99=310.58ms
Single I/O Operation (Per Block):
Read: Avg=14.572ms, P50=0.280ms, P95=0.280ms, P99=0.280ms
Write: Avg=5.120ms, P50=5.120ms, P95=5.120ms, P99=5.120ms
[I/O & Bandwidth]
Total Read I/O: 473.0 MB (473 ops)
Total Write I/O: 1476.0 MB (1,476 ops)
Effective Bandwidth: 280.8 MB/s
[Storage Details]
Blocks in Use: 1,476
Free Blocks: 0
Tokens per Block: 512
Block Size: 1.00 MB
[Execution Time]
Total Execution Time: 8.42 s
================================================================================
```
## Metrics
| Metric | Description |
|--------|-------------|
| **QPS** | Queries per second (based on I/O time, excluding sleep) |
| **Request Latency** | End-to-end latency for entire request (all I/O operations) |
| **Single I/O Latency** | Latency for individual block read/write operations (512 tokens) |
| **P50/P95/P99** | Latency percentiles (milliseconds) using linear interpolation |
| **Hit Rate** | Cache hit ratio for blocks |
| **Write Ratio** | Percentage of blocks that needed to be written |
| **Bandwidth** | Effective throughput based on I/O time only |
| **Prefix Hits** | Number of blocks served from prefix cache |
**Note**: Request Latency measures the total time to process all blocks in a request, while Single I/O Latency measures the time for one block operation (512 tokens).
## Trace Data Format
```json
{
"timestamp": 1234.567,
"hash_ids": [1, 2, 4, 7],
"input_length": 2048,
"output_length": 512
}
```
Each `hash_id` corresponds to a 512-token block. The tool simulates prefix caching by checking if blocks are already in storage before writing.
## Requirements
- Python 3.10+

View File

@ -264,151 +264,6 @@ def get_into(self, key: str, buffer_ptr: int, size: int) -> int
**Returns:** Number of bytes read, or negative on error
#### get_into_ranges()
Retrieve multiple byte ranges from multiple objects into registered buffers (zero-copy).
```python
def get_into_ranges(self, buffer_ptrs: List[int], all_keys: List[List[str]], all_dst_offsets: List[List[List[int]]], all_src_offsets: List[List[List[int]]], all_sizes: List[List[List[int]]]) -> List[List[List[int]]]
```
This API is **buffer-major** and supports **multiple fragments per key**.
Think of the input shape as:
- `buffer_ptrs[i]`: the `i`-th destination buffer
- `all_keys[i][j]`: the `j`-th key that writes into buffer `i`
- `all_dst_offsets[i][j][k]`: destination offset of fragment `k` for key `j` in buffer `i`
- `all_src_offsets[i][j][k]`: source offset of fragment `k` inside key `j` for buffer `i`
- `all_sizes[i][j][k]`: byte size of fragment `k`
For each triple `(i, j, k)`, Mooncake reads the source range
`[all_src_offsets[i][j][k], all_src_offsets[i][j][k] + all_sizes[i][j][k])`
from object `all_keys[i][j]`, then writes it into destination buffer
`buffer_ptrs[i]` at offset `all_dst_offsets[i][j][k]`.
This lets one buffer gather interleaved fragments from multiple keys, and lets one key contribute multiple disjoint fragments to the same buffer in a single call.
**Parameters:**
- `buffer_ptrs`: Memory addresses of pre-allocated destination buffers. Every buffer must be registered with `register_buffer()` before calling this API.
- `all_keys`: For each buffer, the ordered list of source object keys to read from.
- `all_dst_offsets`: For each buffer and key, the destination offsets of that key's fragments.
- `all_src_offsets`: For each buffer and key, the source offsets of that key's fragments inside the object.
- `all_sizes`: For each buffer and key, the byte lengths of that key's fragments.
**Shape rules:**
- `len(buffer_ptrs) == len(all_keys) == len(all_dst_offsets) == len(all_src_offsets) == len(all_sizes)`
- For each buffer `i`, `len(all_keys[i]) == len(all_dst_offsets[i]) == len(all_src_offsets[i]) == len(all_sizes[i])`
- For each `(buffer i, key j)`, `len(all_dst_offsets[i][j]) == len(all_src_offsets[i][j]) == len(all_sizes[i][j])`
If a top-level shape or per-key fragment shape does not match, the corresponding result entries are negative error codes.
**Returns:** A nested list of per-buffer, per-key, per-fragment results. `results[i][j][k]` is the number of bytes read for fragment `k`, or a negative value on error.
A successful call can still contain per-fragment failures. For example, if one key is missing but another key in the same buffer is valid, the missing key's fragment result will be negative while the valid fragment can still succeed.
**Typical scenarios:**
- **Partial read from one object:** You only need a slice of a large value, such as a header, metadata block, or a small subrange of a tensor shard. In this case, use one buffer, one key, and one or more fragments under that key.
- **Stitch multiple fragments from one object into one buffer:** You need several non-contiguous ranges from the same object and want to pack them into one destination buffer. In this case, keep a single key entry and place multiple fragments under that key.
- **Stitch data from multiple objects into one buffer:** You want to assemble one logical payload from several keys. In this case, use one destination buffer and list multiple keys under that buffer, with each key contributing one or more fragments.
- **Fill multiple output buffers in one call:** You have several destination buffers, each with its own read plan. In this case, each top-level entry in `buffer_ptrs` and the parallel nested arrays describes one independent destination buffer.
**How to use it for partial reads:**
If you only want part of an object, do not call `get_into()` with the full object buffer size. Instead:
1. Allocate and register a destination buffer sized for the bytes you actually want to materialize.
2. Put that buffer pointer into `buffer_ptrs`.
3. Put the source key into `all_keys`.
4. Set `all_src_offsets` to the start offsets of the object ranges you want.
5. Set `all_sizes` to the lengths of those ranges.
6. Set `all_dst_offsets` to where those ranges should land in your destination buffer.
A useful way to think about the arguments is:
- `buffer_ptrs` answers **where does the data land**
- `all_keys` answers **which object does it come from**
- `all_src_offsets` and `all_sizes` answer **which bytes should be read**
- `all_dst_offsets` answers **where each fragment should be placed in the destination buffer**
If you are extracting a single contiguous slice from one object, the minimal shape is:
```python
results = store.get_into_ranges(
[buffer_ptr],
[["my_key"]],
[[[0]]],
[[[src_offset]]],
[[[size]]],
)
```
This means:
- one destination buffer
- one source key for that buffer
- one fragment for that key
- read `size` bytes from `my_key[src_offset:src_offset + size]`
- write them into `buffer_ptr[0:size]`
If you want to read several disjoint ranges from the same object and pack them together, keep the same key and add more fragments under it. For example:
```python
results = store.get_into_ranges(
[buffer_ptr],
[["my_key"]],
[[[0, 16, 40]]],
[[[128, 4096, 8192]]],
[[[8, 12, 4]]],
)
```
This reads three fragments from `my_key` and places them into the same destination buffer at offsets `0`, `16`, and `40`. This pattern is useful when you want to assemble only the needed pieces of a large object without reading the whole value.
If you want to assemble one output buffer from multiple objects, keep one top-level buffer entry and add multiple keys under it. Each key can still contribute one or more fragments. For example, you might put a header from `meta_key` at the front of the buffer, then place a payload slice from `data_key` after it.
**Usage example:**
```python
import ctypes
buffer_size = 32
buffer0 = (ctypes.c_ubyte * buffer_size)()
buffer1 = (ctypes.c_ubyte * buffer_size)()
buffer_ptr0 = ctypes.addressof(buffer0)
buffer_ptr1 = ctypes.addressof(buffer1)
store.register_buffer(buffer_ptr0, buffer_size)
store.register_buffer(buffer_ptr1, buffer_size)
# Buffer 0 reads:
# - from key1: two fragments -> src[1:5] -> dst[0:4], src[30:33] -> dst[20:23]
# - from key2: one fragment -> src[2:7] -> dst[8:13]
# Buffer 1 reads:
# - from key2: one fragment -> src[0:6] -> dst[4:10]
# - from key1: one fragment -> src[10:14] -> dst[16:20]
results = store.get_into_ranges(
[buffer_ptr0, buffer_ptr1],
[["key1", "key2"], ["key2", "key1"]],
[[[0, 20], [8]], [[4], [16]]],
[[[1, 30], [2]], [[0], [10]]],
[[[4, 3], [5]], [[6], [4]]],
)
# results == [
# [[4, 3], [5]],
# [[6], [4]],
# ]
```
In the example above:
- `results[0][0][0] == 4`: buffer 0, key 0 (`"key1"`), fragment 0 succeeded with 4 bytes
- `results[0][0][1] == 3`: buffer 0, key 0 (`"key1"`), fragment 1 succeeded with 3 bytes
- `results[0][1][0] == 5`: buffer 0, key 1 (`"key2"`), fragment 0 succeeded with 5 bytes
**Common pitfalls:**
- Do not flatten all fragments for a buffer into one list. Fragments must be grouped under their corresponding key.
- `all_dst_offsets`, `all_src_offsets`, and `all_sizes` are 3D, but `all_keys` is 2D.
- Buffer overflow is checked against the registered destination buffer size.
- Source overflow is checked against the source object's size.
- Full-object `get_into()` and ranged `get_into_ranges()` are different APIs; use `get_into()` when you want the whole object into one buffer.
**Current limitation:** true ranged items currently require the selected source replica to be memory-backed. Whole-object reads still follow the normal full-read path, but partial reads through `get_into_ranges()` do not support non-memory replicas.
---
## ReplicateConfig Configuration
@ -446,16 +301,6 @@ config = ReplicateConfig()
config.with_soft_pin = True # Keep this object in memory longer
```
#### with_hard_pin
**Type:** `bool`
**Default:** `False`
**Description:** Enables hard pinning for the stored object. Hard pinned objects will not be evicted. This grants user to manually control the life time of stored objects.
```python
config = ReplicateConfig()
config.with_hard_pin = True # Keep this object in memory that will not be evicted
```
#### preferred_segment
**Type:** `str`
**Default:** `""` (empty string)
@ -784,120 +629,6 @@ result = store.put_batch(keys, values)
---
#### upsert()
Insert a new object if the key does not exist, or update the existing object in place when possible. They use the same replication configuration model as `put()`.
Upsert binary data in the distributed storage.
```python
def upsert(self, key: str, value: bytes, config: ReplicateConfig = None) -> int
```
**Parameters:**
- `key` (str): Unique object identifier
- `value` (bytes): Binary data to insert or update
- `config` (ReplicateConfig, optional): Replication configuration
**Returns:**
- `int`: Status code (0 = success, non-zero = error code)
**Example:**
```python
config = ReplicateConfig()
config.replica_num = 2
rc = store.upsert("weights", b"new-bytes", config)
if rc == 0:
print("Upsert succeeded")
```
#### upsert_from()
Upsert object data directly from a pre-allocated buffer (zero-copy).
```python
def upsert_from(self, key: str, buffer_ptr: int, size: int, config: ReplicateConfig = None) -> int
```
**Parameters:**
- `key` (str): Object identifier
- `buffer_ptr` (int): Memory address of the source buffer
- `size` (int): Number of bytes to insert or update
- `config` (ReplicateConfig, optional): Replication configuration
**Returns:**
- `int`: Status code (0 = success, non-zero = error code)
**Note:** This is the zero-copy counterpart of `upsert()`. As with
`put_from()`, register the buffer before issuing the request.
#### batch_upsert_from()
Upsert multiple objects directly from pre-allocated buffers.
```python
def batch_upsert_from(self, keys: List[str], buffer_ptrs: List[int], sizes: List[int],
config: ReplicateConfig = None) -> List[int]
```
**Parameters:**
- `keys` (List[str]): List of object identifiers
- `buffer_ptrs` (List[int]): List of source buffer addresses
- `sizes` (List[int]): List of byte lengths for each buffer
- `config` (ReplicateConfig, optional): Replication configuration shared by all objects
**Returns:**
- `List[int]`: List of status codes for each upsert
#### upsert_parts()
Upsert data from multiple buffer parts as a single object (insert or update).
```python
def upsert_parts(self, key: str, *parts, config: ReplicateConfig = None) -> int
```
**Parameters:**
- `key` (str): Object identifier
- `*parts`: Variable number of bytes-like objects to concatenate
- `config` (ReplicateConfig, optional): Replication configuration
**Returns:**
- `int`: Status code (0 = success, non-zero = error code)
**Example:**
```python
part1 = b"Hello, "
part2 = b"World!"
result = store.upsert_parts("greeting", part1, part2)
```
#### upsert_batch()
Upsert multiple objects in a single batch operation.
```python
def upsert_batch(self, keys: List[str], values: List[bytes], config: ReplicateConfig = None) -> int
```
**Parameters:**
- `keys` (List[str]): List of object identifiers
- `values` (List[bytes]): List of binary data to insert or update
- `config` (ReplicateConfig, optional): Replication configuration for all objects
**Returns:**
- `int`: Status code (0 = success, non-zero = error code)
**Example:**
```python
keys = ["key1", "key2", "key3"]
values = [b"value1", b"value2", b"value3"]
result = store.upsert_batch(keys, values)
```
---
#### get_batch()
Retrieve multiple objects in a single batch operation.
@ -990,39 +721,6 @@ print(f"Removed {count} objects")
---
#### batch_remove()
Remove multiple objects by their keys in a single batch operation.
```python
def batch_remove(self, keys: List[str], force: bool = False) -> List[int]
```
**Parameters:**
- `keys` (List[str]): List of object identifiers to remove
- `force` (bool): If True, skip lease and replication task checks (default: False)
**Returns:**
- `List[int]`: List of status codes for each key (0 = success, negative = error code)
**Example:**
```python
# Remove multiple keys in one batch
keys = ["key1", "key2", "key3", "key4", "key5"]
results = store.batch_remove(keys)
# Check results
for key, result in zip(keys, results):
if result == 0:
print(f"✓ {key} removed successfully")
else:
print(f"✗ {key} failed with error code: {result}")
# Force remove (bypass lease checks)
results = store.batch_remove(keys, force=True)
```
---
#### is_exist()
Check if an object exists in the storage system.
@ -1802,133 +1500,6 @@ def batch_pub_tensor(self, keys: List[str], tensors_list: List[torch.Tensor], co
---
#### upsert_tensor()
Insert a tensor if its key is missing, or update the existing tensor if the key already exists. The current tensor upsert helpers use the default `ReplicateConfig` and therefore do not take a `config` parameter.
Upsert a PyTorch tensor into the store.
```python
def upsert_tensor(self, key: str, tensor: torch.Tensor) -> int
```
**Parameters:**
- `key` (str): Object identifier
- `tensor` (torch.Tensor): The PyTorch tensor to insert or update
**Returns:**
- `int`: Status code (0 = success, non-zero = error code)
**Note:** This function requires `torch` to be installed and available in the environment.
#### upsert_tensor_from()
Upsert a tensor directly from a pre-allocated buffer. The buffer layout must be
`[TensorMetadata][tensor data]`, matching the layout used by
`get_tensor_into()`.
```python
def upsert_tensor_from(self, key: str, buffer_ptr: int, size: int) -> int
```
**Parameters:**
- `key` (str): Object identifier
- `buffer_ptr` (int): Buffer pointer containing serialized tensor metadata and payload
- `size` (int): Actual serialized byte length of the tensor buffer
**Returns:**
- `int`: Status code (0 = success, non-zero = error code)
**Note:** This function is not supported for dummy client.
#### batch_upsert_tensor_from()
Upsert multiple tensors directly from pre-allocated buffers. Each buffer must
use layout `[TensorMetadata][tensor data]`.
```python
def batch_upsert_tensor_from(self, keys: List[str], buffer_ptrs: List[int], sizes: List[int]) -> List[int]
```
**Parameters:**
- `keys` (List[str]): List of object identifiers
- `buffer_ptrs` (List[int]): List of serialized tensor buffer pointers
- `sizes` (List[int]): List of actual serialized byte lengths
**Returns:**
- `List[int]`: List of status codes for each tensor upsert
#### batch_upsert_tensor()
Upsert a batch of PyTorch tensors into the store (insert or update).
```python
def batch_upsert_tensor(self, keys: List[str], tensors_list: List[torch.Tensor]) -> List[int]
```
**Parameters:**
- `keys` (List[str]): List of object identifiers
- `tensors_list` (List[torch.Tensor]): List of tensors to insert or update
**Returns:**
- `List[int]`: List of status codes for each tensor operation.
**Note:** This function requires `torch` to be installed and available in the environment. Not supported for dummy client.
#### upsert_pub_tensor()
Upsert a PyTorch tensor with configurable replication settings (insert or update).
```python
def upsert_pub_tensor(self, key: str, tensor: torch.Tensor, config: ReplicateConfig = None) -> int
```
**Parameters:**
- `key` (str): Unique object identifier
- `tensor` (torch.Tensor): PyTorch tensor to insert or update
- `config` (ReplicateConfig, optional): Replication configuration
**Returns:**
- `int`: Status code (0 = success, non-zero = error code)
**Note:** This function requires `torch` to be installed and available in the environment. Not supported for dummy client.
**Example:**
```python
import torch
from mooncake.store import ReplicateConfig
tensor = torch.randn(100, 100)
config = ReplicateConfig()
config.replica_num = 2
config.with_soft_pin = True
result = store.upsert_pub_tensor("my_tensor", tensor, config)
if result == 0:
print("Tensor upserted successfully")
```
#### batch_upsert_pub_tensor()
Batch upsert PyTorch tensors with configurable replication settings (insert or update).
```python
def batch_upsert_pub_tensor(self, keys: List[str], tensors_list: List[torch.Tensor], config: ReplicateConfig = None) -> List[int]
```
**Parameters:**
- `keys` (List[str]): List of object identifiers
- `tensors_list` (List[torch.Tensor]): List of tensors to insert or update
- `config` (ReplicateConfig, optional): Replication configuration
**Returns:**
- `List[int]`: List of status codes for each tensor operation.
**Note:** This function requires `torch` to be installed and available in the environment. Not supported for dummy client.
---
### PyTorch Tensor Operations (Zero Copy)
These methods provide direct support for storing and retrieving PyTorch tensors. They automatically handle serialization and metadata, and include built-in support for **Tensor Parallelism (TP)** by automatically splitting and reconstructing tensor shards.
@ -1964,8 +1535,8 @@ def batch_get_tensor_into(self, base_keys: List[str], buffer_ptrs: List[int], si
**Parameters:**
- `base_keys` (List[str]): List of base identifiers.
- `buffer_ptrs` (List[int]): List of buffer pointers pre-allocated for tensor; buffers should be registered.
- `sizes` (List[int]): List of buffer sizes.
- `buffer_ptrs` (List[int]): List of the buffers pointer pre-allocated for tensor, and the buffers should be registered.
- `sizes` (List[int]): List of the size of buffers.
**Returns:**
@ -2003,8 +1574,8 @@ def batch_get_tensor_with_tp_into(self, base_keys: List[str], buffer_ptrs: List[
**Parameters:**
- `base_keys` (List[str]): List of base identifiers.
- `buffer_ptrs` (List[int]): List of buffer pointers pre-allocated for tensor; buffers should be registered.
- `sizes` (List[int]): List of buffer sizes.
- `buffer_ptrs` (List[int]): List of the buffers pointer pre-allocated for tensor, and the buffers should be registered.
- `sizes` (List[int]): List of the size of buffers.
- `tp_rank` (int): The tensor parallel rank to retrieve (default: 0).
- `tp_size` (int): Total tensor parallel size (default: 1).
@ -2012,84 +1583,6 @@ def batch_get_tensor_with_tp_into(self, base_keys: List[str], buffer_ptrs: List[
- `List[torch.Tensor]`: List of retrieved tensors (or shards). Contains `None` for missing keys.
#### put_tensor_from()
Put a PyTorch tensor into the store directly from a pre-allocated buffer (zero-copy). The buffer must contain data in the same layout as produced by `get_tensor_into`: **\[TensorMetadata\]\[tensor data\]**. The buffer is only read during this call; no Python object references it.
```python
def put_tensor_from(self, key: str, buffer_ptr: int, size: int) -> int
```
**Parameters:**
- `key` (str): Object identifier for the tensor.
- `buffer_ptr` (int): The buffer pointer; the buffer should be registered. Layout must be \[TensorMetadata\]\[tensor data\].
- `size` (int): **Actual serialized byte length** of the data in the buffer (metadata + tensor bytes), not the buffer capacity.
**Returns:**
- `int`: Status code (0 = success, non-zero = error code).
#### batch_put_tensor_from()
Put a batch of PyTorch tensors into the store directly from pre-allocated buffers (zero-copy). Each buffer must contain data in the layout **\[TensorMetadata\]\[tensor data\]**, same as `get_tensor_into`.
```python
def batch_put_tensor_from(self, keys: List[str], buffer_ptrs: List[int], sizes: List[int]) -> List[int]
```
**Parameters:**
- `keys` (List[str]): List of object identifiers.
- `buffer_ptrs` (List[int]): List of buffer pointers; buffers should be registered.
- `sizes` (List[int]): List of **actual serialized byte lengths** for each buffer (metadata + tensor bytes), not buffer capacities.
**Returns:**
- `List[int]`: List of status codes for each tensor operation (0 = success, non-zero = error code).
#### put_tensor_with_tp_from()
Put a **full tensor** into the store directly from a pre-allocated buffer (zero-copy), for use with Tensor Parallelism. This is the zero-copy counterpart of `put_tensor_with_tp()`: the buffer must contain the complete tensor in layout **\[TensorMetadata\]\[tensor data\]**, and Mooncake will split it internally and store all shards under `key_tp_<rank>`.
```python
def put_tensor_with_tp_from(self, key: str, buffer_ptr: int, size: int, tp_rank: int = 0, tp_size: int = 1, split_dim: int = 0) -> int
```
**Parameters:**
- `key` (str): Base identifier for the tensor.
- `buffer_ptr` (int): The buffer pointer; the buffer should be registered.
- `size` (int): **Actual serialized byte length** of the full tensor in the buffer.
- `tp_rank` (int): Kept for signature compatibility with `put_tensor_with_tp()` (default: 0). It does **not** mean "only write one shard".
- `tp_size` (int): Total tensor parallel size (default: 1). If 1, equivalent to `put_tensor_from(key, buffer_ptr, size)`.
- `split_dim` (int): Dimension along which the full tensor is split before storing shards.
**Returns:**
- `int`: Status code (0 = success, non-zero = error code).
#### batch_put_tensor_with_tp_from()
Put a batch of **full tensors** into the store directly from pre-allocated buffers (zero-copy). This is the zero-copy counterpart of `batch_put_tensor_with_tp()`: each buffer contains one full tensor in layout **\[TensorMetadata\]\[tensor data\]**, and Mooncake splits each tensor internally and stores all TP shards.
```python
def batch_put_tensor_with_tp_from(self, base_keys: List[str], buffer_ptrs: List[int], sizes: List[int], tp_rank: int = 0, tp_size: int = 1, split_dim: int = 0) -> List[int]
```
**Parameters:**
- `base_keys` (List[str]): List of base identifiers.
- `buffer_ptrs` (List[int]): List of buffer pointers; buffers should be registered.
- `sizes` (List[int]): List of **actual serialized byte lengths** for each full-tensor buffer.
- `tp_rank` (int): Kept for signature compatibility with `batch_put_tensor_with_tp()` (default: 0). It does **not** select a single shard to write.
- `tp_size` (int): Total tensor parallel size (default: 1). If 1, equivalent to `batch_put_tensor_from(base_keys, buffer_ptrs, sizes)`.
- `split_dim` (int): Dimension along which each full tensor is split before storing shards.
**Returns:**
- `List[int]`: List of status codes for each tensor operation (0 = success, non-zero = error code).
---
### Batch Zero-Copy Operations

View File

@ -63,37 +63,7 @@ Errors in this part usually indicate that the error occurred within the `mooncak
**Solution:**
Ensure that the total memory registration does not exceed the device's upper limit. You may need to reduce the amount of memory being registered or split large memory regions into smaller chunks that fit within the device's `max_mr_size` limit.
5. If you encounter `Failed to register memory 0x...: Resource temporarily unavailable [11]` and kernel logs show `CREATE_MKEY failed, status no resources(0xf)`, this indicates that the RDMA NIC has exhausted its internal Memory Key (MKEY) resources, even though `ulimit -l` and `vm.max_map_count` may appear sufficient.
This typically happens when:
- Applications that use RDMA (e.g., SGLang with HiCache + Mooncake) have crashed or been killed multiple times without cleanly releasing RDMA resources.
- The leaked MKEY entries accumulate in the NIC firmware and are not reclaimed by the kernel, eventually hitting the hardware limit.
- Large memory regions (e.g., NSA indexer buffers at ~4.68 GB each across multiple TP ranks) amplify the problem since each registration consumes more internal NIC resources.
**Diagnostic Commands:**
```bash
# Check current RDMA resource usage per device
rdma resource show
# Check kernel logs for CREATE_MKEY failures
dmesg | grep -i "CREATE_MKEY\|no resources\|mlx5_cmd_out_err"
# Example output:
# mlx5_core 0000:65:01.0: mlx5_cmd_out_err:829:(pid 3958462): CREATE_MKEY(0x200) op_mod(0x0) failed, status no resources(0xf), syndrome (0x2aac7c), err(-11)
# Ensure vm.max_map_count is large enough (default 65530 may be too small)
sysctl vm.max_map_count
```
**Solutions:**
- **Reboot the node** to fully reset NIC firmware state and reclaim all leaked MKEY resources. This is the most reliable fix.
- Increase `vm.max_map_count` if it is at the default value: `sysctl -w vm.max_map_count=16777216`
- Ensure applications shut down cleanly (avoid `kill -9` when possible) so RDMA resources are properly deregistered.
- If rebooting is not feasible, try unloading and reloading the mlx5 kernel modules (may disrupt other services):
```bash
modprobe -r mlx5_ib mlx5_core && modprobe mlx5_core mlx5_ib
```
6. If you encounter errors indicating inability to allocate memory space when requesting large memory regions, this may be due to ulimit restrictions. When the total memory requirement (number of registered RDMA devices × requested space) exceeds the ulimit, the system will display errors about failing to allocate space.
5. If you encounter errors indicating inability to allocate memory space when requesting large memory regions, this may be due to ulimit restrictions. When the total memory requirement (number of registered RDMA devices × requested space) exceeds the ulimit, the system will display errors about failing to allocate space.
**Diagnostic Commands:**
- Use `ulimit -a` to check current limits, particularly the `max locked memory` value
@ -110,7 +80,7 @@ Errors in this part usually indicate that the error occurred within the `mooncak
* hard memlock unlimited
```
7. If the error `Failed to create QP: Cannot allocate memory` is displayed, it is typically caused by too many QP have been created, reaching the driver limit. You can use `rdma resource` to trace how many QP is created. One possible way to resolve this issue:
6. If the error `Failed to create QP: Cannot allocate memory` is displayed, it is typically caused by too many QP have been created, reaching the driver limit. You can use `rdma resource` to trace how many QP is created. One possible way to resolve this issue:
- Update Mooncake to version v0.3.5 or later
- Set the environment variable `MC_ENABLE_DEST_DEVICE_AFFINITY=1` before starting the application

View File

@ -1,8 +1,5 @@
# Ascend Transport
Ascend Transport源代码路径为Mooncake/mooncake-transfer-engine/src/transport/ascend_transport该路径下还包含自动化编译脚本、README文件。
**Ascend Transport 已不再维护,昇腾平台推荐使用 [Ascend Direct Transport](./ascend_direct_transport.md). **
## 概述
Ascend Transport是一个单边语义的高性能零拷贝NPU数据传输库直接兼容Mooncake Transfer Engine。要编译使用Ascend Transport库请在mooncake-common\common.cmake文件中将USE_ASCEND开关置于"ON"。

View File

@ -110,41 +110,7 @@
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/musa/lib
```
4. 若需编译寒武纪 MLU 支持,请先安装寒武纪 Neuware SDK。之后
1) 导出 `NEUWARE_HOME`,或在 CMake 中传入 `-DNEUWARE_ROOT=/path/to/neuware`
2) 配置 `LIBRARY_PATH``LD_LIBRARY_PATH`,确保编译时能链接 `cnrt`、`cndrv` 等 Neuware 库:
```bash
export NEUWARE_HOME=/usr/local/neuware
export LIBRARY_PATH=$LIBRARY_PATH:${NEUWARE_HOME}/lib64
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${NEUWARE_HOME}/lib64
```
若 Neuware 安装路径与默认头文件/库布局不一致,还可显式指定:
```bash
cmake .. -DUSE_MLU=ON \
-DMLU_INCLUDE_DIR=/path/to/neuware/include \
-DMLU_LIB_DIR=/path/to/neuware/lib64
```
启用 MLU 后端示例:
```bash
cmake .. -DUSE_MLU=ON -DNEUWARE_ROOT=${NEUWARE_HOME:-/usr/local/neuware}
make -j
```
5. 若需编译沐曦 MetaX MACA 支持(如 C500请安装 MACA SDK使头文件与库位于 `MACA_ROOT`(优先取 `MACA_HOME` 环境变量,未设置时默认 `/opt/maca`)。不同安装包可能把库放在 `lib``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 列表)
6. 安装 yalantinglibs
4. 安装 yalantinglibs
```bash
git clone https://github.com/alibaba/yalantinglibs.git
cd yalantinglibs
@ -154,7 +120,7 @@
make install
```
7. 进入项目根目录,运行下列命令进行编译
5. 进入项目根目录,运行下列命令进行编译
```bash
mkdir build
cd build
@ -162,7 +128,7 @@
make -j
```
8. 安装 Mooncake python 包和 mooncake_master 可执行文件
6. 安装 Mooncake python 包和 mooncake_master 可执行文件
```bash
make install
```
@ -171,14 +137,6 @@
在执行 `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_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 头文件与库目录。
- `-DUSE_HIP=[ON|OFF]`: 通过 HIP/ROCm 启用对 AMD GPU 的支持
- `-DUSE_CXL=[ON|OFF]`: 启用 CXL 支持
- `-DWITH_STORE=[ON|OFF]`: 编译 Mooncake Store 组件

View File

@ -1,90 +0,0 @@
# Kunpeng UB Transport
Kunpeng UbTransport源代码路径为Mooncake/mooncake-transfer-engine/src/transport/kunpneg_transport该路径下有UB协议的Transport对接代码和实现逻辑。
## 概述
UBUnified Bus统一总线 是与RDMA、CXL、NVLink 和TCP处于同一抽象层的传输协议属于可在应用层灵活选择的传输方案。目前 UB 协议有两个开源实现URMA远程内存访问语义和 OBMMLoad/Store 语义)。
URMAUnified Remote Memory Access统一远程内存访问是UB协议为上层应用提供的统一编程抽象与核心语义层。它基于 UB 协议低延迟、高带宽的底层特性,为远程共享内存的访问与操作提供统一的 API 和语义接口。
URMA 开源代码仓库https://atomgit.com/openeuler/umdk
OBMM (Ownership Based Memory Management) 是面向超节点环境的内核内存管理系统,支持跨节点的物理内存共享。该系统通过内核模块 (obmm.ko) 和用户态库 (libobmm.so) 提供高效的远程内存访问能力。
OBMM 开源代码仓库https://atomgit.com/openeuler/obmm
## 新增依赖
Kunpeng UbTransport在Mooncake本身依赖的基础上新增了一部分URMA和OBMM的依赖
- **硬件平台**: 支持原生UB互联架构的鲲鹏950 CPU
- **OS版本**: openEuler 24.03 (LTS-SP3) [下载链接](https://www.openeuler.openatom.cn/zh/download/#openEuler%2024.03%20LTS%20SP3)
- **URMA依赖**: UMDK: `yum install umdk-urma-devel` 或从[源码](https://atomgit.com/openeuler/umdk)构建。
- **协议优势**: URMA 提供类似 RDMA 的内存语义,针对鲲鹏芯片片上互联进行了优化
---
## 构建与编译
**前置条件**
- openEuler 24.03 (LTS-SP3) [下载链接](https://www.openeuler.openatom.cn/zh/download/#openEuler%2024.03%20LTS%20SP3)
- 已安装 UMDK: `yum install umdk-urma-devel` 或从[源码](https://atomgit.com/openeuler/umdk)构建
**CMake 配置**
```bash
# 克隆 Mooncake 仓库
git clone https://github.com/kvcache-ai/Mooncake.git
cd Mooncake
# 启用 UB 传输层进行配置
mkdir build && cd build
cmake .. -DUSE_UB=ON \
-DURMA_INCLUDE_DIR=/usr/include \
-DURMA_LIBRARY=/usr/lib64/liburma.so
# 编译
make -j$(nproc)
```
**验证**
```bash
# 检查 UB 传输层是否已注册
./mooncake_server --list-transports
# 预期输出: rdma, tcp, nvlink, ub
```
---
## 运行与测试
**单节点基准测试**
```bash
# 终端 1: 目标端Target
./transfer_engine_bench \
--mode=target \
--protocol=ub \
--device_name=urma0 \
--local_server_name=127.0.0.1 \
--metadata_server=P2PHANDSHAKE
# 终端 2: 发起端Initiator
./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
```
**多设备基准测试**
```bash
# 自动发现多个 URMA 设备
./transfer_engine_bench \
--protocol=ub \
--device_name=urma0,urma1,urma2,urma3
```

View File

@ -779,8 +779,6 @@ Max threads: 4
Master service listening on 0.0.0.0:50051
```
如果 Master 运行在容器中,而容器 IP 可能动态变化,建议使用 `--rpc-interface=<网卡名>`(例如 `--rpc-interface=eth0`)而不是写死 `--rpc-address`。Master 会在启动时解析该网卡当前的 IPv4 地址,并将其作为最终的 `rpc_address` 使用。
**高可用模式**:
高可用模式依赖于 etcd 服务进行协调。如果 Transfer Engine 也使用 etcd 作为其元数据服务,那么 Mooncake Store 使用的 etcd 集群可以与 Transfer Engine 使用的集群共用,也可以是独立的。
@ -790,7 +788,6 @@ Master service listening on 0.0.0.0:50051
--enable-ha启用高可用模式
--etcd-endpoints指定 etcd 服务的多个入口,使用分号 ';' 分隔
--rpc-address该实例的 RPC 地址。注意,这里填写的地址应当是客户端可访问的地址。
--rpc-interface按网卡名解析当前实例的 IPv4 地址。设置后会覆盖 --rpc-address适合容器 IP 会变化的场景。
```
例如:
@ -801,14 +798,6 @@ Master service listening on 0.0.0.0:50051
--rpc-address=10.0.0.1
```
容器部署示例:
```
./build/mooncake-store/src/mooncake_master \
--enable-ha=true \
--etcd-endpoints="0.0.0.0:2379;0.0.0.0:2479;0.0.0.0:2579" \
--rpc-interface=eth0
```
### 启动验证程序
Mooncake Store 提供了多种验证程序,包括基于 C++ 和 Python 等接口形态。下面以 `stress_cluster_benchmark` 为例介绍一下如何运行。

@ -1 +0,0 @@
Subproject commit 73dea196d23ad8fcd4914c6ef1238f390b9a1c48

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

@ -1,64 +0,0 @@
# SetupPyTorchEnv.cmake
#
# This file provides helper functions for building Mooncake Pytorch extensions
# and is meant to be included by BuildEpExt.cmake and BuildPgExt.cmake.
# Ensure we have the correct Python interpreter (respects active virtualenvs)
find_package(Python3 REQUIRED COMPONENTS Interpreter)
# Install PyTorch for a specific version with proper CUDA compatibility handling.
#
# Usage:
# install_pytorch_wheel("<VERSION>" <CUDA_MAJOR> <CUDA_MINOR> "<MODULE_PREFIX>")
#
# Example:
# install_pytorch_wheel("2.11.0" 12 8 "[EP]")
function(install_pytorch_wheel _version _cuda_major _cuda_minor _module_prefix)
message(STATUS "${_module_prefix} Installing PyTorch ${_version} via pip...")
set(_cu_tag "")
# Determine the specific CUDA tag for PyTorch wheels
if(_cuda_major GREATER_EQUAL 13)
# TODO: Fix when we need to support more CUDA 13 versions or when the CI env is fixed.
set(_cu_tag "cu130")
elseif(_cuda_major EQUAL 12 AND _version VERSION_GREATER_EQUAL "2.11.0")
# PyTorch 2.11.0+ defaults to CUDA 13.
# We must explicitly point to CUDA 12 wheels for these newer versions.
if(_cuda_minor GREATER_EQUAL 8)
set(_cu_tag "cu128")
elseif(_cuda_minor GREATER_EQUAL 6)
set(_cu_tag "cu126")
else()
message(FATAL_ERROR
"${_module_prefix} Can't find a matching PyTorch wheel for version ${_version} "
"with CUDA ${_cuda_major}.${_cuda_minor}"
)
endif()
endif()
# Construct pip command using the absolute path to the Python executable
set(_pip_cmd ${Python3_EXECUTABLE} -m pip install "torch==${_version}")
if(_cu_tag)
set(_index_url "https://download.pytorch.org/whl/${_cu_tag}")
message(STATUS "${_module_prefix} Using specific CUDA wheel: ${_index_url}")
list(APPEND _pip_cmd --index-url "${_index_url}")
else()
message(STATUS "${_module_prefix} Using default PyPI wheels for PyTorch ${_version}")
endif()
# Execute pip install
execute_process(
COMMAND ${_pip_cmd}
RESULT_VARIABLE _ret
)
if(NOT _ret EQUAL 0)
message(FATAL_ERROR "${_module_prefix} Failed to install PyTorch ${_version}."
" Command run: '${_pip_cmd}'")
endif()
message(STATUS "${_module_prefix} PyTorch ${_version} is ready.")
endfunction()

View File

@ -1,10 +0,0 @@
# SetupPython.cmake resolve the Python interpreter for execute_process() calls.
#
# Honour -DPython3_EXECUTABLE=... when provided (e.g. Docker builds that
# install a non-system Python via deadsnakes), otherwise fall back to the
# default "python3" on PATH. Sets PYTHON_EXECUTABLE for legacy callers.
if(NOT Python3_EXECUTABLE)
set(Python3_EXECUTABLE "python3")
endif()
set(PYTHON_EXECUTABLE "${Python3_EXECUTABLE}")

View File

@ -40,9 +40,6 @@ add_definitions(-DCONFIG_ERDMA)
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
# Memory-aware build parallelism (compile vs. link job pools)
include(${CMAKE_CURRENT_LIST_DIR}/limit_jobs.cmake)
option(ENABLE_SCCACHE "Whether to open sccache" OFF)
if (ENABLE_SCCACHE)
find_program(SCCACHE sccache REQUIRED)
@ -60,9 +57,7 @@ option(BUILD_EXAMPLES "Build examples" ON)
option(BUILD_UNIT_TESTS "Build unit tests" ON)
option(USE_CUDA "option for enabling gpu features for NVIDIA GPU" OFF)
option(USE_MLU "option for enabling Cambricon MLU features" OFF)
option(USE_MUSA "option for enabling gpu features for MTHREADS GPU" OFF)
option(USE_MACA "option for enabling gpu features for MUXI GPU with MACA" OFF)
option(USE_HIP "option for enabling gpu features for AMD GPU" OFF)
option(USE_NVMEOF "option for using NVMe over Fabric" OFF)
option(USE_TCP "option for using TCP transport" ON)
@ -74,13 +69,6 @@ option(USE_ASCEND_HETEROGENEOUS "option for transferring between ascend npu and
option(USE_MNNVL "option for using Multi-Node NVLink transport" OFF)
option(USE_CXL "option for using CXL protocol" OFF)
option(USE_EFA "option for using AWS EFA transport" OFF)
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
@ -114,11 +102,7 @@ option(WITH_NVIDIA_PEERMEM "disable to support RDMA without nvidia-peermem. If W
option(USE_EVENT_DRIVEN_COMPLETION "option for using event-driven completion (store & transfer engine)" OFF)
option(USE_TENT "option for building Mooncake TENT" OFF)
option(ENABLE_MULTI_PROTOCOL "option for enabling multi-protocol support in transfer engine" OFF)
if (ENABLE_MULTI_PROTOCOL)
add_compile_definitions(ENABLE_MULTI_PROTOCOL)
message(STATUS "Multi-protocol support is enabled")
endif()
option(USE_LRU_MASTER "option for using LRU in master service" OFF)
option(USE_INTRA_NVLINK "option for using IntraNode nvlink transport" OFF)
set(LRU_MAX_CAPACITY 1000)
@ -142,7 +126,7 @@ if (USE_NVMEOF)
endif()
if (USE_MNNVL)
if (NOT USE_HIP AND NOT USE_MUSA AND NOT USE_MACA)
if (NOT USE_HIP AND NOT USE_MUSA)
set(USE_CUDA ON)
endif()
add_compile_definitions(USE_MNNVL)
@ -159,60 +143,6 @@ if (USE_CUDA)
)
endif()
if (NOT DEFINED NEUWARE_ROOT OR NEUWARE_ROOT STREQUAL "")
if (DEFINED ENV{NEUWARE_HOME} AND NOT "$ENV{NEUWARE_HOME}" STREQUAL "")
set(NEUWARE_ROOT "$ENV{NEUWARE_HOME}" CACHE PATH "Path to Cambricon Neuware SDK" FORCE)
else()
set(NEUWARE_ROOT "/usr/local/neuware" CACHE PATH "Path to Cambricon Neuware SDK" FORCE)
endif()
endif()
if (NOT DEFINED MLU_INCLUDE_DIR OR MLU_INCLUDE_DIR STREQUAL "")
set(MLU_INCLUDE_DIR "${NEUWARE_ROOT}/include")
endif()
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")
include_directories(${MLU_INCLUDE_DIR})
if (EXISTS "${MLU_LIB_DIR}")
link_directories(${MLU_LIB_DIR})
endif()
endif()
if (USE_MACA)
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})
endif()
endif()
if (USE_MUSA)
add_compile_definitions(USE_MUSA)
message(STATUS "MUSA support is enabled")

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

@ -34,13 +34,11 @@ import "C"
import (
"context"
"encoding/json"
"errors"
"strings"
"sync"
"time"
"unsafe"
rpctypes "go.etcd.io/etcd/api/v3/v3rpc/rpctypes"
clientv3 "go.etcd.io/etcd/client/v3"
)
@ -66,11 +64,11 @@ var (
storeKeepAliveCtx = make(map[int64]context.CancelFunc)
storeKeepAliveMutex sync.Mutex
// watch contexts for store
storeWatchCtx = make(map[string]context.CancelFunc)
storeWatchMutex sync.Mutex
storeWatchCtx = make(map[string]context.CancelFunc)
storeWatchMutex sync.Mutex
// etcd client for HA snapshot
snapshotClient *clientv3.Client
snapshotMutex sync.Mutex
snapshotClient *clientv3.Client
snapshotMutex sync.Mutex
// watch contexts for prefix watch
storePrefixWatchCtx = make(map[string]prefixWatchInfo)
storePrefixWatchMutex sync.Mutex
@ -78,7 +76,7 @@ var (
const (
// Snapshot client config (for GB-level snapshot files)
snapshotMaxMsgSize = 2000 * 1000 * 1000 // 2GB
snapshotMaxMsgSize = 2000 * 1000 * 1000 // 2GB
snapshotTimeout = 60 * time.Second // 1 minute for large files
)
@ -327,25 +325,6 @@ func EtcdStoreGrantLeaseWrapper(ttl int64, leaseId *int64, errMsg **C.char) int
return 0
}
//export EtcdStoreRevokeLeaseWrapper
func EtcdStoreRevokeLeaseWrapper(leaseId int64, errMsg **C.char) int {
if storeClient == nil {
*errMsg = C.CString("etcd client not initialized")
return -1
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
_, err := storeClient.Revoke(ctx, clientv3.LeaseID(leaseId))
if err != nil {
if errors.Is(err, rpctypes.ErrLeaseNotFound) {
return 0
}
*errMsg = C.CString(err.Error())
return -1
}
return 0
}
//export EtcdStoreCreateWithLeaseWrapper
func EtcdStoreCreateWithLeaseWrapper(key *C.char, keySize C.int, value *C.char, valueSize C.int,
leaseId int64, revisionId *int64, errMsg **C.char) int {
@ -480,14 +459,6 @@ func cancelAndDeleteKeepAlive(leaseId int64) int {
return -1
}
func hasKeepAliveContext(leaseId int64) bool {
storeKeepAliveMutex.Lock()
defer storeKeepAliveMutex.Unlock()
_, exists := storeKeepAliveCtx[leaseId]
return exists
}
//export EtcdStoreKeepAliveWrapper
func EtcdStoreKeepAliveWrapper(leaseId int64, errMsg **C.char) int {
if storeClient == nil {
@ -547,21 +518,6 @@ func EtcdStoreCancelKeepAliveWrapper(leaseId int64, errMsg **C.char) int {
return 0
}
//export EtcdStoreWaitKeepAliveReadyWrapper
func EtcdStoreWaitKeepAliveReadyWrapper(leaseId int64, timeoutMs int, errMsg **C.char) int {
deadline := time.Now().Add(time.Duration(timeoutMs) * time.Millisecond)
for {
if hasKeepAliveContext(leaseId) {
return 0
}
if timeoutMs <= 0 || !time.Now().Before(deadline) {
*errMsg = C.CString("keep alive context did not become ready before timeout")
return -1
}
time.Sleep(time.Millisecond)
}
}
//export EtcdStorePutWrapper
func EtcdStorePutWrapper(key *C.char, keySize C.int, value *C.char, valueSize C.int, errMsg **C.char) int {
if storeClient == nil {

View File

@ -1,28 +1,27 @@
module github.com/kvcache-ai/Mooncake/mooncake-common/etcd
go 1.25.0
go 1.23.0
toolchain go1.25.9
toolchain go1.23.7
require (
go.etcd.io/etcd/api/v3 v3.5.21
go.etcd.io/etcd/client/v3 v3.5.21
)
require go.etcd.io/etcd/client/v3 v3.5.21
require (
github.com/coreos/go-semver v0.3.0 // indirect
github.com/coreos/go-systemd/v22 v22.3.2 // indirect
github.com/gogo/protobuf v1.3.2 // indirect
github.com/golang/protobuf v1.5.4 // indirect
go.etcd.io/etcd/api/v3 v3.5.21 // indirect
go.etcd.io/etcd/client/pkg/v3 v3.5.21 // indirect
go.uber.org/atomic v1.7.0 // indirect
go.uber.org/multierr v1.6.0 // indirect
go.uber.org/zap v1.17.0 // indirect
golang.org/x/net v0.48.0 // indirect
golang.org/x/sys v0.39.0 // indirect
golang.org/x/text v0.32.0 // indirect
google.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217 // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217 // indirect
google.golang.org/grpc v1.79.3 // indirect
google.golang.org/protobuf v1.36.10 // indirect
golang.org/x/net v0.38.0 // indirect
golang.org/x/sys v0.31.0 // indirect
golang.org/x/text v0.23.0 // indirect
google.golang.org/genproto v0.0.0-20230822172742-b8732ec3820d // indirect
google.golang.org/genproto/googleapis/api v0.0.0-20230822172742-b8732ec3820d // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20230822172742-b8732ec3820d // indirect
google.golang.org/grpc v1.59.0 // indirect
google.golang.org/protobuf v1.33.0 // indirect
)

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

@ -70,18 +70,6 @@ class DefaultConfig {
void GetUInt64(const std::string& key, uint64_t* val,
uint64_t default_value = 0) const;
/**
* @brief GetDurationMs retrieves a duration value from the configuration
* and converts it to milliseconds.
* @param key The key to look up in the configuration
* @param val Pointer to store the retrieved value in milliseconds
* @param default_value Default value to return if the key is not found
* @note Duration strings may use ms, s, m, or h as suffixes. Bare numbers
* are interpreted as milliseconds.
*/
void GetDurationMs(const std::string& key, uint64_t* val,
uint64_t default_value = 0) const;
/**
* @brief GetDouble retrieves a double value from the configuration
* @param key The key to look up in the configuration

View File

@ -1,95 +0,0 @@
#pragma once
#include <cctype>
#include <cstdint>
#include <limits>
#include <string>
#include <string_view>
namespace mooncake {
inline std::string_view TrimAsciiWhitespace(std::string_view value) {
while (!value.empty() &&
std::isspace(static_cast<unsigned char>(value.front()))) {
value.remove_prefix(1);
}
while (!value.empty() &&
std::isspace(static_cast<unsigned char>(value.back()))) {
value.remove_suffix(1);
}
return value;
}
inline bool ParseDurationMs(std::string_view value, uint64_t* result,
std::string* error = nullptr) {
auto set_error = [&](std::string message) {
if (error != nullptr) {
*error = std::move(message);
}
return false;
};
if (result == nullptr) {
return set_error("duration output pointer is null");
}
std::string_view trimmed = TrimAsciiWhitespace(value);
if (trimmed.empty()) {
return set_error(
"duration is empty; expected a non-negative integer optionally "
"followed by ms, s, m, or h");
}
size_t number_end = 0;
while (number_end < trimmed.size() &&
std::isdigit(static_cast<unsigned char>(trimmed[number_end]))) {
++number_end;
}
if (number_end == 0) {
return set_error(
"duration must start with a non-negative integer and may use ms, "
"s, m, or h as the unit suffix");
}
uint64_t numeric_value = 0;
for (size_t i = 0; i < number_end; ++i) {
const uint64_t digit = static_cast<uint64_t>(trimmed[i] - '0');
if (numeric_value >
(std::numeric_limits<uint64_t>::max() - digit) / 10) {
return set_error("duration value is too large");
}
numeric_value = numeric_value * 10 + digit;
}
std::string_view suffix = TrimAsciiWhitespace(trimmed.substr(number_end));
std::string normalized_suffix;
normalized_suffix.reserve(suffix.size());
for (char ch : suffix) {
normalized_suffix.push_back(
static_cast<char>(std::tolower(static_cast<unsigned char>(ch))));
}
uint64_t multiplier = 1;
if (normalized_suffix.empty() || normalized_suffix == "ms") {
multiplier = 1;
} else if (normalized_suffix == "s") {
multiplier = 1000;
} else if (normalized_suffix == "m") {
multiplier = 60 * 1000;
} else if (normalized_suffix == "h") {
multiplier = 60 * 60 * 1000;
} else {
return set_error("unsupported duration unit '" + normalized_suffix +
"'; supported units are ms, s, m, and h");
}
if (numeric_value > std::numeric_limits<uint64_t>::max() / multiplier) {
return set_error("duration value is too large after unit conversion");
}
*result = numeric_value * multiplier;
return true;
}
} // namespace mooncake

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,86 +0,0 @@
# limit_jobs.cmake Memory-aware build parallelism
#
# Auto-detects available memory and CPU count, calculates safe parallel job
# limits for compilation and linking separately. With Ninja, creates job pools
# so compilation uses many cores while memory-heavy linking is restricted.
#
# User overrides (cmake -D...):
# PARALLEL_COMPILE_JOBS override compile parallelism
# PARALLEL_LINK_JOBS override link parallelism
# MAX_COMPILER_MEMORY_MB per-compile-job memory estimate (default: 1500)
# MAX_LINKER_MEMORY_MB per-link-job memory estimate (default: 4000)
set(MAX_COMPILER_MEMORY_MB "1500" CACHE STRING
"Estimated peak memory per compile job in MB")
set(MAX_LINKER_MEMORY_MB "4000" CACHE STRING
"Estimated peak memory per link job in MB")
# Guard against invalid user input (division by zero)
if(MAX_COMPILER_MEMORY_MB LESS_EQUAL 0)
message(WARNING "[limit_jobs] MAX_COMPILER_MEMORY_MB=${MAX_COMPILER_MEMORY_MB} "
"invalid, falling back to 1500")
set(MAX_COMPILER_MEMORY_MB 1500 CACHE STRING
"Estimated peak memory per compile job in MB" FORCE)
endif()
if(MAX_LINKER_MEMORY_MB LESS_EQUAL 0)
message(WARNING "[limit_jobs] MAX_LINKER_MEMORY_MB=${MAX_LINKER_MEMORY_MB} "
"invalid, falling back to 4000")
set(MAX_LINKER_MEMORY_MB 4000 CACHE STRING
"Estimated peak memory per link job in MB" FORCE)
endif()
# Detect system resources
cmake_host_system_information(RESULT _available_mem_mb
QUERY AVAILABLE_PHYSICAL_MEMORY)
cmake_host_system_information(RESULT _nproc
QUERY NUMBER_OF_LOGICAL_CORES)
message(STATUS "[limit_jobs] Available memory: ${_available_mem_mb} MB, "
"CPU cores: ${_nproc}")
# Calculate safe parallel jobs from memory
math(EXPR _compile_jobs "${_available_mem_mb} / ${MAX_COMPILER_MEMORY_MB}")
math(EXPR _link_jobs "${_available_mem_mb} / ${MAX_LINKER_MEMORY_MB}")
# Clamp: [1, nproc]
if(_compile_jobs LESS 1)
set(_compile_jobs 1)
endif()
if(_compile_jobs GREATER _nproc)
set(_compile_jobs ${_nproc})
endif()
if(_link_jobs LESS 1)
set(_link_jobs 1)
endif()
if(_link_jobs GREATER _nproc)
set(_link_jobs ${_nproc})
endif()
# Use auto-detected values unless user explicitly overrides with -D
if(NOT DEFINED PARALLEL_COMPILE_JOBS)
set(PARALLEL_COMPILE_JOBS "${_compile_jobs}")
endif()
if(NOT DEFINED PARALLEL_LINK_JOBS)
set(PARALLEL_LINK_JOBS "${_link_jobs}")
endif()
message(STATUS "[limit_jobs] Compile jobs: ${PARALLEL_COMPILE_JOBS} "
"(${MAX_COMPILER_MEMORY_MB} MB/job), "
"Link jobs: ${PARALLEL_LINK_JOBS} (${MAX_LINKER_MEMORY_MB} MB/job)")
# Apply to build system
if(CMAKE_GENERATOR MATCHES "Ninja")
set_property(GLOBAL APPEND PROPERTY JOB_POOLS
compile_pool=${PARALLEL_COMPILE_JOBS}
link_pool=${PARALLEL_LINK_JOBS}
)
set(CMAKE_JOB_POOL_COMPILE "compile_pool" CACHE STRING "" FORCE)
set(CMAKE_JOB_POOL_LINK "link_pool" CACHE STRING "" FORCE)
message(STATUS "[limit_jobs] Ninja job pools: "
"compile=${PARALLEL_COMPILE_JOBS}, link=${PARALLEL_LINK_JOBS}")
else()
message(STATUS "[limit_jobs] Hint: use -G Ninja for automatic "
"compile/link parallelism separation")
message(STATUS "[limit_jobs] With Make, recommend: "
"cmake --build . -j${PARALLEL_LINK_JOBS}")
endif()

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

@ -1,7 +1,5 @@
#include "default_config.h"
#include "duration_utils.h"
#if __has_include(<jsoncpp/json/reader.h>)
#include <jsoncpp/json/reader.h>
#include <jsoncpp/json/value.h> // Ubuntu
@ -164,71 +162,6 @@ void DefaultConfig::GetUInt64(const std::string& key, uint64_t* val,
}
}
void DefaultConfig::GetDurationMs(const std::string& key, uint64_t* val,
uint64_t default_value) const {
Node node;
if (!getValue(key, &node)) {
*val = default_value;
return;
}
if (type_ == ConfigType::YAML) {
std::string raw_value = node.yaml_node_.as<std::string>();
std::string error;
if (!ParseDurationMs(raw_value, val, &error)) {
throw std::runtime_error("Invalid duration for key '" + key +
"': " + error);
}
return;
}
if (type_ == ConfigType::JSON) {
if (node.json_value_.isString()) {
std::string error;
if (!ParseDurationMs(node.json_value_.asString(), val, &error)) {
throw std::runtime_error("Invalid duration for key '" + key +
"': " + error);
}
return;
}
if (node.json_value_.isUInt64()) {
*val = node.json_value_.asUInt64();
return;
}
if (node.json_value_.isUInt()) {
*val = static_cast<uint64_t>(node.json_value_.asUInt());
return;
}
if (node.json_value_.isInt64()) {
const int64_t numeric_value = node.json_value_.asInt64();
if (numeric_value < 0) {
throw std::runtime_error("Invalid duration for key '" + key +
"': value must be non-negative");
}
*val = static_cast<uint64_t>(numeric_value);
return;
}
if (node.json_value_.isInt()) {
const int numeric_value = node.json_value_.asInt();
if (numeric_value < 0) {
throw std::runtime_error("Invalid duration for key '" + key +
"': value must be non-negative");
}
*val = static_cast<uint64_t>(numeric_value);
return;
}
throw std::runtime_error("Invalid duration for key '" + key +
"': JSON value must be an integer or string");
}
*val = default_value;
}
void DefaultConfig::GetDouble(const std::string& key, double* val,
double default_value) const {
Node node;

View File

@ -1,5 +1,4 @@
#include "default_config.h"
#include "duration_utils.h"
#include <gtest/gtest.h>
@ -78,93 +77,6 @@ TEST_F(DefaultConfigTest, LoadYamlSuccess) {
ASSERT_EQ(config_data_.default_value, 10000);
};
TEST(DurationUtilsTest, ParseDurationMsSupportsLegacyMillisecondsAndUnits) {
uint64_t value = 0;
ASSERT_TRUE(ParseDurationMs("5000", &value));
ASSERT_EQ(value, 5000);
ASSERT_TRUE(ParseDurationMs("5000ms", &value));
ASSERT_EQ(value, 5000);
ASSERT_TRUE(ParseDurationMs("5s", &value));
ASSERT_EQ(value, 5000);
ASSERT_TRUE(ParseDurationMs("30m", &value));
ASSERT_EQ(value, 30 * 60 * 1000);
ASSERT_TRUE(ParseDurationMs("1H", &value));
ASSERT_EQ(value, 60 * 60 * 1000);
ASSERT_TRUE(ParseDurationMs(" 7 m ", &value));
ASSERT_EQ(value, 7 * 60 * 1000);
}
TEST(DurationUtilsTest, ParseDurationMsRejectsInvalidInput) {
uint64_t value = 0;
ASSERT_FALSE(ParseDurationMs("", &value));
ASSERT_FALSE(ParseDurationMs("abc", &value));
ASSERT_FALSE(ParseDurationMs("-1", &value));
ASSERT_FALSE(ParseDurationMs("1d", &value));
ASSERT_FALSE(ParseDurationMs("18446744073709551616", &value));
ASSERT_FALSE(ParseDurationMs("18446744073709552h", &value));
}
TEST_F(DefaultConfigTest, GetDurationMsFromJsonSupportsNumbersAndStrings) {
DefaultConfig config;
config.SetPath(path_ + "/../../mooncake-common/tests/test.json");
config.Load();
uint64_t legacy_ms = 0;
uint64_t seconds = 0;
uint64_t minutes = 0;
uint64_t hours = 0;
uint64_t whitespace = 0;
uint64_t missing_default = 0;
config.GetDurationMs("legacyDurationMs", &legacy_ms, 0);
config.GetDurationMs("durationSeconds", &seconds, 0);
config.GetDurationMs("durationMinutes", &minutes, 0);
config.GetDurationMs("durationHours", &hours, 0);
config.GetDurationMs("durationWhitespace", &whitespace, 0);
config.GetDurationMs("missingDuration", &missing_default, 1234);
ASSERT_EQ(legacy_ms, 5000);
ASSERT_EQ(seconds, 5000);
ASSERT_EQ(minutes, 30 * 60 * 1000);
ASSERT_EQ(hours, 60 * 60 * 1000);
ASSERT_EQ(whitespace, 7 * 60 * 1000);
ASSERT_EQ(missing_default, 1234);
}
TEST_F(DefaultConfigTest, GetDurationMsFromYamlSupportsNumbersAndStrings) {
DefaultConfig config;
config.SetPath(path_ + "/../../mooncake-common/tests/test.yaml");
config.Load();
uint64_t legacy_ms = 0;
uint64_t seconds = 0;
uint64_t minutes = 0;
uint64_t hours = 0;
uint64_t whitespace = 0;
uint64_t missing_default = 0;
config.GetDurationMs("legacyDurationMs", &legacy_ms, 0);
config.GetDurationMs("durationSeconds", &seconds, 0);
config.GetDurationMs("durationMinutes", &minutes, 0);
config.GetDurationMs("durationHours", &hours, 0);
config.GetDurationMs("durationWhitespace", &whitespace, 0);
config.GetDurationMs("missingDuration", &missing_default, 4321);
ASSERT_EQ(legacy_ms, 6000);
ASSERT_EQ(seconds, 6000);
ASSERT_EQ(minutes, 7 * 60 * 1000);
ASSERT_EQ(hours, 60 * 60 * 1000);
ASSERT_EQ(whitespace, 8 * 60 * 60 * 1000);
ASSERT_EQ(missing_default, 4321);
}
TEST_F(DefaultConfigTest, LoadInvalidFile) {
DefaultConfig config;
config.SetPath(path_ + "/invalid_file.txt");
@ -175,4 +87,4 @@ TEST_F(DefaultConfigTest, LoadInvalidFile) {
int main(int argc, char **argv) {
::testing::InitGoogleTest(&argc, argv);
return RUN_ALL_TESTS();
}
}

View File

@ -3,13 +3,8 @@
"testFloat": 4.15,
"testString": "Hello, World",
"testBoolean": false,
"legacyDurationMs": 5000,
"durationSeconds": "5s",
"durationMinutes": "30m",
"durationHours": "1H",
"durationWhitespace": " 7 m ",
"testObject": {
"nestedInteger": 1000,
"nestedString": "Nested Hello, World"
}
}
}

View File

@ -2,11 +2,6 @@ testInteger: 42
testFloat: 3.14
testString: "Hello, World!"
testBoolean: true
legacyDurationMs: 6000
durationSeconds: 6s
durationMinutes: "7m"
durationHours: 1H
durationWhitespace: " 8 h "
testObject:
nestedInteger: 100
nestedString: "Nested Hello"

View File

@ -0,0 +1,17 @@
find_program(GO_EXECUTABLE go)
if(NOT GO_EXECUTABLE)
message(FATAL_ERROR "Go compiler not found. Please install Golang first.")
endif()
add_custom_target(mooncake_conductor ALL
COMMAND ./build.sh ${CMAKE_CURRENT_BINARY_DIR} ${CMAKE_BINARY_DIR}
WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
COMMENT "Building Go program: mooncake_conductor"
VERBATIM
)
set(GO_EXECUTABLE_PATH "${CMAKE_CURRENT_BINARY_DIR}/mooncake_conductor"
CACHE INTERNAL "Path to Go executable")
# Install target
install(PROGRAMS ${GO_EXECUTABLE_PATH} DESTINATION bin)

View File

@ -0,0 +1,167 @@
# vLLM V1 Disaggregated Serving with MooncakeConductor
## Overview
This is the latest version of the Mooncake Conductor integration doc with the vLLM project to support KVCache-Aware scheduling algorithm.
The conductor can be integrated as a plugin into any proxy to uniformly manage KV events from L1 to L3. We also provide a toy_proxy for those who want to try it out ([proxy](./cacheaware_disaggregated_proxy.py)). Benchmark results will be released soon.
- only vLLM and vLLM-Ascend are supported.
## Installation
The mooncake conductor will be compiled and installed together with the mooncake store. Refer to [Build Guide](https://github.com/kvcache-ai/Mooncake/blob/main/docs/source/getting_started/build.md).
- **WITH_CONDUCTOR must be set to ON in Mooncake/CMakeLists.txt.**
### Install the latest version of vLLM and vLLM-Ascend
#### 1. Clone vLLM from official repo
```bash
git clone git@github.com:vllm-project/vllm.git
```
#### 2. Build
##### 2.1 Build from source
```bash
cd vllm
pip3 install -e .
```
- If you encounter any problems that you cannot solve, please refer to the [vLLM official compilation guide](https://docs.vllm.ai/en/latest/getting_started/installation/index.html).
#### 3. Clone vLLM-Ascend from official repo
```bash
git clone git@github.com:vllm-project/vllm-ascend.git
```
#### 4. Build
##### 4.1 Build from source
```bash
cd vllm-ascend
pip install -e .
```
- If you encounter any problems that you cannot solve, please refer to the [vLLM-Ascend official compilation guide](https://docs.vllm.ai/projects/ascend/en/latest/).
## Configuration
### Prepare configuration file to Run Example
- Prepare a _**conductor_config.json**_ file for mooncake_conductor. Here is an example:
```json
{
"kvevent_instance":
{
"vllm-1":
{
"ip": "127.0.0.1",
"port": 5557,
"type": "vLLM",
"modelname": "qwen2.5_7B",
"lora_id": -1
},
"mooncake":
{
"ip": "127.0.0.1",
"port": 19997,
"type": "Mooncake",
"modelname": "qwen2.5_7B",
"lora_id": -1
}
},
"http_server_port": 13333
}
```
- `kvevent_instance`: Services capable of reporting KV events.
- `vllm-1/mooncake`: rename of a VLLM instance or Mooncake-master instance.You can modify it according to your own preferences.
- `ip`: zmq publisher IP.
- `port`: zmq publisher port.
- `type`: Mark the type of kv-event publisher. Generally, there are currently only two types: `vLLM` and `Mooncake`.
- `modelname`: Model name used for match the model.
- `lora_id`: LoRA Adapter ID.
- `http_server_port`: Conductor http server for querying cache hit rates, default use `13333`.
## Run Example
### 1. Start the mooncake_master server
```sh
# start mooncake_master without kv-event publish
mooncake_master --rpc_port 50051
# start moocake_master with kv-event
mooncake_master -enable_kv_event_publish -kv_event_publisher_endpoint tcp://*:19997 -rpc_port 50051
```
### 2. Run multiple vllm instances
```sh
# kv_producer role
vllm serve /qwen2.5_7B_instruct/ \
--enforce-eager \
--max-model-len 10000 \
--port 8100 \
--gpu-memory-utilization 0.8 \
--served-model-name "qwen2.5_7B" \
--trust-remote-code \
--kv-events-config \
'{
"publisher": "zmq",
"enable_kv_cache_events": true,
"endpoint": "tcp://*:5557",
"topic": "kv-events",
"replay_endpoint": "tcp://*:5558"
}' \
--kv-transfer-config \
'{
"kv_connector": "MooncakeConnectorStoreV1",
"kv_role":"kv_producer",
"kv_connector_extra_config":{"use_layerwise": false}
}'
```
```sh
# kv_consumer role
vllm serve /qwen2.5_7B_instruct/ \
--enforce-eager \
--max-model-len 10000 \
--port 8200 \
--gpu-memory-utilization 0.8 \
--served-model-name "qwen2.5_7B" \
--trust-remote-code \
--kv-transfer-config \
'{
"kv_connector": "MooncakeConnectorStoreV1",
"kv_role":"kv_consumer",
"kv_connector_extra_config":{"use_layerwise": false}
}'
```
### 3. Start the conductor server
```sh
export CONDUCTOR_CONFIG_PATH="./example/conductor_config.json"
mooncake_conductor
```
### 4. Run the proxy in the example
```sh
python cacheaware_disaggregated_proxy.py --prefiller-hosts 127.0.0.1 --prefiller-ports 8100 --decoder-host 127.0.0.1 --decoder-ports 8200 --conductor-address 127.0.0.1:13333
```
## Test with openai compatible request
```sh
curl -s http://localhost:8000/v1/completions -H "Content-Type: application/json" -d '{
"model": "qwen2.5_7B",
"prompt": "What are the key architectural differences between vLLM and Mooncake when it comes to handling key-value (KV) cache events, and how can a centralized conductor component be designed in Go to normalize disparate event schemas from these systems, apply consistent metrics collection, and make dynamic scheduling decisions based on real-time KV cache hit rates without relying on Kubernetes-based autoscaling mechanisms?",
"max_tokens": 1000
}'
```

33
mooncake-conductor/build.sh Executable file
View File

@ -0,0 +1,33 @@
#!/bin/bash
if [ "$#" -ne 2 ]; then
echo "Usage: $0 TARGET_PATH BUILD_DIR"
exit 1
fi
TARGET=$1
BUILD_DIR=$2
cd conductor-ctrl
# Check if go.mod exists
if [ ! -f "go.mod" ]; then
echo "Error: go.mod file not found"
exit 1
fi
echo "Cleaning previous build..."
rm -f mooncake_conductor
go mod tidy
echo "Building Go program: mooncake_conductor"
go build -o "$TARGET/mooncake_conductor" main.go
if [ $? -eq 0 ] && [ -f "$TARGET/mooncake_conductor" ]; then
echo "mooncake_conductor built successfully"
else
echo "mooncake_conductor build failed"
exit 1
fi

View File

@ -0,0 +1,82 @@
package common
import (
"sync"
"sync/atomic"
)
type SyncMap[K any, V any] struct {
m sync.Map
len atomic.Int32
}
func (sm *SyncMap[K, V]) Delete(key K) {
sm.LoadAndDelete(key)
}
func (sm *SyncMap[K, V]) Load(key K) (typedVal V, ok bool) {
value, ok := sm.m.Load(key)
if ok {
typedVal = value.(V)
}
return
}
func (sm *SyncMap[K, V]) LoadAndDelete(key K) (typedVal V, loaded bool) {
value, loaded := sm.m.LoadAndDelete(key)
if loaded {
typedVal = value.(V)
sm.len.Add(-1)
}
return
}
func (sm *SyncMap[K, V]) LoadOrStore(key K, value V) (V, bool) {
actual, loaded := sm.m.LoadOrStore(key, value)
if !loaded {
sm.len.Add(1)
}
return actual.(V), loaded
}
func (sm *SyncMap[K, V]) Range(f func(key K, value V) bool) {
sm.m.Range(func(key, value any) bool {
return f(key.(K), value.(V))
})
}
func (sm *SyncMap[K, V]) Keys() []K {
k := make([]K, 0, sm.Len())
sm.m.Range(func(key, value any) bool {
k = append(k, key.(K))
return true
})
return k
}
func (sm *SyncMap[K, V]) Values() []V {
v := make([]V, 0, sm.Len())
sm.m.Range(func(key, value any) bool {
v = append(v, value.(V))
return true
})
return v
}
func (sm *SyncMap[K, V]) Store(key K, value V) {
sm.Swap(key, value)
}
func (sm *SyncMap[K, V]) Swap(key K, value V) (V, bool) {
old, loaded := sm.m.Swap(key, value)
if !loaded {
var ret V
sm.len.Add(1)
return ret, loaded
}
return old.(V), loaded
}
func (sm *SyncMap[K, V]) Len() int {
return int(sm.len.Load())
}

View File

@ -0,0 +1,145 @@
package common
import (
"fmt"
"sync"
"testing"
)
func TestSyncMap(t *testing.T) {
sm := &SyncMap[string, int]{}
// test Store and Load
sm.Store("key1", 1)
val, ok := sm.Load("key1")
if !ok || val != 1 {
t.Errorf("Expected Load('key1') to return (1, true), got (%d, %v)", val, ok)
}
// test Len
if sm.Len() != 1 {
t.Errorf("Expected Len() to return 1, got %d", sm.Len())
}
// test Load non-existent key
val, ok = sm.Load("key2")
if ok {
t.Errorf("Expected Load('key2') to return (0, false), got (%d, %v)", val, ok)
}
// test LoadOrStore with non-existent key
val, loaded := sm.LoadOrStore("key2", 2)
if loaded || val != 2 {
t.Errorf("Expected LoadOrStore('key2', 2) to return (2, false), got (%d, %v)", val, loaded)
}
if sm.Len() != 2 {
t.Errorf("Expected Len() to return 2, got %d", sm.Len())
}
// test LoadOrStore with existing key
val, loaded = sm.LoadOrStore("key1", 10)
if !loaded || val != 1 {
t.Errorf("Expected LoadOrStore('key1', 10) to return (1, true), got (%d, %v)", val, loaded)
}
// test Swap with existing key
val, loaded = sm.Swap("key1", 3)
if !loaded || val != 1 {
t.Errorf("Expected Swap('key1', 3) to return (1, true), got (%d, %v)", val, loaded)
}
val, ok = sm.Load("key1")
if !ok || val != 3 {
t.Errorf("Expected Load('key1') after Swap to return (3, true), got (%d, %v)", val, ok)
}
// test Swap with non-existent key
val, loaded = sm.Swap("key3", 4)
if loaded || val != 0 {
t.Errorf("Expected Swap('key3', 4) to return (0, false), got (%d, %v)", val, loaded)
}
if sm.Len() != 3 {
t.Errorf("Expected Len() to return 3, got %d", sm.Len())
}
// test Keys
keys := sm.Keys()
if len(keys) != 3 {
t.Errorf("Expected Keys() to return 3 keys, got %d", len(keys))
}
// test Keys is exists
keyMap := make(map[string]bool)
for _, k := range keys {
keyMap[k] = true
}
expectedKeys := []string{"key1", "key2", "key3"}
for _, k := range expectedKeys {
if !keyMap[k] {
t.Errorf("Expected key '%s' in Keys() result", k)
}
}
values := sm.Values()
if len(values) != 3 {
t.Errorf("Expected Values() to return 3 values, got %d", len(values))
}
// test Range
var rangeCount int
sm.Range(func(key string, value int) bool {
rangeCount++
return true
})
if rangeCount != 3 {
t.Errorf("Expected Range to iterate over 3 items, got %d", rangeCount)
}
// test Range with early termination
var earlyTerminateCount int
sm.Range(func(key string, value int) bool {
earlyTerminateCount++
return earlyTerminateCount < 2 // only iterate over the first two elements
})
if earlyTerminateCount != 2 {
t.Errorf("Expected Range with early termination to iterate over 2 items, got %d", earlyTerminateCount)
}
// test LoadAndDelete
val, loaded = sm.LoadAndDelete("key2")
if !loaded || val != 2 {
t.Errorf("Expected LoadAndDelete('key2') to return (2, true), got (%d, %v)", val, loaded)
}
if sm.Len() != 2 {
t.Errorf("Expected Len() after LoadAndDelete to return 2, got %d", sm.Len())
}
// test Delete
sm.Delete("key1")
if sm.Len() != 1 {
t.Errorf("Expected Len() after Delete to return 1, got %d", sm.Len())
}
val, ok = sm.Load("key1")
if ok {
t.Errorf("Expected Load('key1') after Delete to return (0, false), got (%d, %v)", val, ok)
}
// test concurrent safety
var wg sync.WaitGroup
concurrency := 100
wg.Add(concurrency)
for i := 0; i < concurrency; i++ {
go func(i int) {
defer wg.Done()
key := fmt.Sprintf("concurrent_key_%d", i)
sm.Store(key, i)
val, ok := sm.Load(key)
if !ok || val != i {
t.Errorf("Concurrent test failed: Expected Load('%s') to return (%d, true), got (%d, %v)", key, i, val, ok)
}
}(i)
}
wg.Wait()
if sm.Len() != concurrency+1 { // +1 for key3 still present
t.Errorf("Expected Len() after concurrent operations to return %d, got %d", concurrency+1, sm.Len())
}
}

View File

@ -0,0 +1,39 @@
package common
const (
ServiceTypeVLLM string = "vLLM"
ServiceTypeMooncake string = "Mooncake"
)
type ServiceConfig struct {
Endpoint string // kv publisher endpoint
ReplayEndpoint string // (optional)
Type string // kv publisher type, support: vLLM,Mooncake
ModelName string // Model name hosted by the service
LoraName string
TenantID string // (optional), default use 'default'
InstanceID string // required
BlockSize int64
DPRank int
AdditionalSalt string // (optional), default use empty string
}
type StoredEvent struct {
BlockHashes []uint64
BlockSize int64
ModelName string
LoraName string
InstanceID string
ParentBlockHash uint64
TokenIds []int32
Medium string
}
type RemovedEvent struct {
BlockHashes []uint64
ModelName string
LoraName string
InstanceID string
BlockSize int64
Medium string
}

View File

@ -0,0 +1,152 @@
package common
import (
"fmt"
"log/slog"
"os"
"strconv"
"strings"
)
func ParseLogLevel() slog.Level {
levelStr := os.Getenv("CONDUCTOR_LOG_LEVEL")
if levelStr == "" {
return slog.LevelInfo
}
switch strings.ToUpper(levelStr) {
case "DEBUG":
return slog.LevelDebug
case "INFO":
return slog.LevelInfo
case "WARN":
return slog.LevelWarn
case "ERROR":
return slog.LevelError
default:
// We use the default logger here to warn about the invalid config
slog.Warn("Invalid log level specified, defaulting to INFO", "level", levelStr)
return slog.LevelInfo
}
}
func LoadEnv(envName, defaultEnv string) string {
value := os.Getenv(envName)
if value == "" {
slog.Warn("environment variable is not set, using default value", "envName", envName, "defaultValue", defaultEnv)
return defaultEnv
}
return value
}
func LoadIntEnv(envName string, defaultEnv int) int {
value := os.Getenv(envName)
trimmedValue := strings.TrimSpace(value)
if value != "" {
intValue, err := strconv.Atoi(value)
if err != nil {
slog.Error("invalid value for environment variable", "envName", envName, "value", trimmedValue)
} else {
return intValue
}
}
slog.Warn("environment variable is not set, using default value", "envName", envName, "defaultValue", defaultEnv)
return defaultEnv
}
func LoadBoolEnv(envName string, defaultEnv bool) bool {
value := os.Getenv(envName)
trimmedValue := strings.TrimSpace(value)
if value != "" {
boolValue, err := strconv.ParseBool(value)
if err != nil {
slog.Error("invalid value for environment variable", "envName", envName, "value", trimmedValue)
} else {
return boolValue
}
}
slog.Warn("environment variable is not set, using default value", "envName", envName, "defaultValue", defaultEnv)
return defaultEnv
}
func LoadFloatEnv(envName string, defaultEnv float64) float64 {
value := os.Getenv(envName)
trimmedValue := strings.TrimSpace(value)
if value != "" {
floatValue, err := strconv.ParseFloat(value, 64)
if err != nil {
slog.Error("invalid value for environment variable", "envName", envName, "value", trimmedValue)
} else {
return floatValue
}
}
slog.Warn("environment variable is not set, using default value", "envName", envName, "defaultValue", defaultEnv)
return defaultEnv
}
func ExtractTokenIdFromRequest(data map[string]interface{}, key string) ([]int32, error) {
raw, exists := data[key]
if !exists {
return nil, fmt.Errorf("missing key: %s", key)
}
arr, ok := raw.([]interface{})
if !ok {
return nil, fmt.Errorf("the value of %s is not an array", key)
}
result := make([]int32, len(arr))
for i, v := range arr {
switch val := v.(type) {
case float64:
result[i] = int32(val)
case int:
result[i] = int32(val)
default:
return nil, fmt.Errorf("unsupported value type of token_id at [%d], the type is %s", i, val)
}
}
return result, nil
}
func ExtractCandidateEngineFromRequest(data map[string]interface{}, key string) (map[string]struct{}, error) {
raw, exists := data[key]
if !exists {
return nil, fmt.Errorf("missing key: %s", key)
}
arr, ok := raw.([]interface{})
if !ok {
return nil, fmt.Errorf("the value of %s is not an array", key)
}
result := make(map[string]struct{}, len(arr))
for i, v := range arr {
str, ok := v.(string)
if !ok {
return nil, fmt.Errorf(`"instances[%d]" is not a string`, i)
}
result[str] = struct{}{}
}
return result, nil
}
func ExtractStringValueFromRequest(data map[string]interface{}, key string) (string, error) {
raw, exists := data[key]
if !exists {
return "", fmt.Errorf("missing key: %s", key)
}
str, ok := raw.(string)
if !ok {
return "", fmt.Errorf(`"the value of: %s" is not a string`, key)
}
return str, nil
}
func ExtractIntFromRequest(data map[string]interface{}, key string) (int64, error) {
raw, exists := data[key]
if !exists {
return -1, fmt.Errorf("missing key: %s", key)
}
result, ok := raw.(float64)
if !ok {
return -1, fmt.Errorf(`"the value of: %s" is not a number`, key)
}
return int64(result), nil
}

View File

@ -0,0 +1,265 @@
package common
import (
"log/slog"
"os"
"testing"
)
func TestParseLogLevel(t *testing.T) {
// test default log level
origLevel := os.Getenv("CONDUCTOR_LOG_LEVEL")
os.Unsetenv("CONDUCTOR_LOG_LEVEL")
defer os.Setenv("CONDUCTOR_LOG_LEVEL", origLevel)
level := ParseLogLevel()
if level != 0 { // slog.LevelInfo = 0
t.Errorf("Expected default log level to be Info, got %d", level)
}
// test various log levels
testCases := []struct {
name string
levelStr string
expected slog.Level
}{
{"Debug", "DEBUG", slog.LevelDebug},
{"Info", "INFO", slog.LevelInfo},
{"Warn", "WARN", slog.LevelWarn},
{"Error", "ERROR", slog.LevelError},
{"Lowercase", "debug", slog.LevelDebug},
{"MixedCase", "Debug", slog.LevelDebug},
{"Invalid", "INVALID", slog.LevelInfo},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
os.Setenv("CONDUCTOR_LOG_LEVEL", tc.levelStr)
level := ParseLogLevel()
if level != tc.expected {
t.Errorf("Expected log level %d for input %s, got %d", tc.expected, tc.levelStr, level)
}
})
}
}
func TestLoadEnv(t *testing.T) {
// test environment variable exists
origValue := os.Getenv("TEST_ENV_VAR")
os.Setenv("TEST_ENV_VAR", "test_value")
defer func() {
if origValue == "" {
os.Unsetenv("TEST_ENV_VAR")
} else {
os.Setenv("TEST_ENV_VAR", origValue)
}
}()
value := LoadEnv("TEST_ENV_VAR", "default_value")
if value != "test_value" {
t.Errorf("Expected LoadEnv to return 'test_value', got '%s'", value)
}
// test environment variable does not exist
origValue2 := os.Getenv("NON_EXISTENT_ENV_VAR")
os.Unsetenv("NON_EXISTENT_ENV_VAR")
defer func() {
if origValue2 != "" {
os.Setenv("NON_EXISTENT_ENV_VAR", origValue2)
}
}()
value = LoadEnv("NON_EXISTENT_ENV_VAR", "default_value")
if value != "default_value" {
t.Errorf("Expected LoadEnv to return 'default_value', got '%s'", value)
}
}
func TestLoadIntEnv(t *testing.T) {
// test environment variable exists and value is valid
origValue := os.Getenv("TEST_INT_ENV_VAR")
os.Setenv("TEST_INT_ENV_VAR", "42")
defer func() {
if origValue == "" {
os.Unsetenv("TEST_INT_ENV_VAR")
} else {
os.Setenv("TEST_INT_ENV_VAR", origValue)
}
}()
value := LoadIntEnv("TEST_INT_ENV_VAR", 100)
if value != 42 {
t.Errorf("Expected LoadIntEnv to return 42, got %d", value)
}
// test environment variable exists but value is invalid
os.Setenv("TEST_INT_ENV_VAR", "invalid")
value = LoadIntEnv("TEST_INT_ENV_VAR", 100)
if value != 100 {
t.Errorf("Expected LoadIntEnv to return 100 for invalid value, got %d", value)
}
// test environment variable does not exist
origValue2 := os.Getenv("NON_EXISTENT_INT_ENV_VAR")
os.Unsetenv("NON_EXISTENT_INT_ENV_VAR")
defer func() {
if origValue2 != "" {
os.Setenv("NON_EXISTENT_INT_ENV_VAR", origValue2)
}
}()
value = LoadIntEnv("NON_EXISTENT_INT_ENV_VAR", 100)
if value != 100 {
t.Errorf("Expected LoadIntEnv to return 100 for non-existent env, got %d", value)
}
}
func TestExtractTokenIdFromRequest(t *testing.T) {
// test successful extraction
data := map[string]interface{}{
"token_ids": []interface{}{1.0, 2.0, 3.0},
}
result, err := ExtractTokenIdFromRequest(data, "token_ids")
if err != nil {
t.Errorf("Expected ExtractTokenIdFromRequest to succeed, got error: %v", err)
}
expected := []int32{1, 2, 3}
if len(result) != len(expected) {
t.Errorf("Expected length %d, got %d", len(expected), len(result))
}
for i, v := range result {
if v != expected[i] {
t.Errorf("Expected result[%d] = %d, got %d", i, expected[i], v)
}
}
// test mixed number types
data["token_ids"] = []interface{}{1, 2.0, 3}
result, err = ExtractTokenIdFromRequest(data, "token_ids")
if err != nil {
t.Errorf("Expected ExtractTokenIdFromRequest to succeed with mixed number types, got error: %v", err)
}
expected = []int32{1, 2, 3}
if len(result) != len(expected) {
t.Errorf("Expected length %d, got %d", len(expected), len(result))
}
for i, v := range result {
if v != expected[i] {
t.Errorf("Expected result[%d] = %d, got %d", i, expected[i], v)
}
}
// test missing key
result, err = ExtractTokenIdFromRequest(data, "missing_key")
if err == nil {
t.Errorf("Expected ExtractTokenIdFromRequest to fail with missing key, got success")
}
// test non-array value
data["token_ids"] = "not an array"
result, err = ExtractTokenIdFromRequest(data, "token_ids")
if err == nil {
t.Errorf("Expected ExtractTokenIdFromRequest to fail with non-array value, got success")
}
// test unsupported types
data["token_ids"] = []interface{}{"string", true}
result, err = ExtractTokenIdFromRequest(data, "token_ids")
if err == nil {
t.Errorf("Expected ExtractTokenIdFromRequest to fail with unsupported types, got success")
}
}
func TestExtractCandidateEngineFromRequest(t *testing.T) {
// test successful extraction
data := map[string]interface{}{
"instances": []interface{}{"engine1", "engine2", "engine3"},
}
result, err := ExtractCandidateEngineFromRequest(data, "instances")
if err != nil {
t.Errorf("Expected ExtractCandidateEngineFromRequest to succeed, got error: %v", err)
}
if len(result) != 3 {
t.Errorf("Expected 3 engines, got %d", len(result))
}
expectedEngines := []string{"engine1", "engine2", "engine3"}
for _, engine := range expectedEngines {
if _, ok := result[engine]; !ok {
t.Errorf("Expected engine '%s' not found in result", engine)
}
}
// test missing key
result, err = ExtractCandidateEngineFromRequest(data, "missing_key")
if err == nil {
t.Errorf("Expected ExtractCandidateEngineFromRequest to fail with missing key, got success")
}
// test non-array value
data["instances"] = "not an array"
result, err = ExtractCandidateEngineFromRequest(data, "instances")
if err == nil {
t.Errorf("Expected ExtractCandidateEngineFromRequest to fail with non-array value, got success")
}
// test non-string element
data["instances"] = []interface{}{"engine1", 2, "engine3"}
result, err = ExtractCandidateEngineFromRequest(data, "instances")
if err == nil {
t.Errorf("Expected ExtractCandidateEngineFromRequest to fail with non-string element, got success")
}
}
func TestExtractStringValueFromRequest(t *testing.T) {
// test successful extraction
data := map[string]interface{}{
"str_key": "string_value",
}
result, err := ExtractStringValueFromRequest(data, "str_key")
if err != nil {
t.Errorf("Expected ExtractStringValueFromRequest to succeed, got error: %v", err)
}
if result != "string_value" {
t.Errorf("Expected 'string_value', got '%s'", result)
}
// test missing key
result, err = ExtractStringValueFromRequest(data, "missing_key")
if err == nil {
t.Errorf("Expected ExtractStringValueFromRequest to fail with missing key, got success")
}
// test non-string value
data["str_key"] = 123
result, err = ExtractStringValueFromRequest(data, "str_key")
if err == nil {
t.Errorf("Expected ExtractStringValueFromRequest to fail with non-string value, got success")
}
}
func TestExtractIntFromRequest(t *testing.T) {
// test successful extraction
data := map[string]interface{}{
"int_key": 42.0,
}
result, err := ExtractIntFromRequest(data, "int_key")
if err != nil {
t.Errorf("Expected ExtractIntFromRequest to succeed, got error: %v", err)
}
if result != 42 {
t.Errorf("Expected 42, got %d", result)
}
// test missing key
result, err = ExtractIntFromRequest(data, "missing_key")
if err == nil {
t.Errorf("Expected ExtractIntFromRequest to fail with missing key, got success")
}
// test non-number value
data["int_key"] = "not a number"
result, err = ExtractIntFromRequest(data, "int_key")
if err == nil {
t.Errorf("Expected ExtractIntFromRequest to fail with non-number value, got success")
}
}

View File

@ -0,0 +1,9 @@
module conductor
go 1.23.8
require (
github.com/cespare/xxhash/v2 v2.3.0
github.com/pebbe/zmq4 v1.4.0
github.com/shamaton/msgpack/v2 v2.4.0
)

View File

@ -0,0 +1,112 @@
package kvevent
import (
"context"
"fmt"
"time"
"conductor/common"
"conductor/zmq"
"log/slog"
)
// KVEventHandler adapts the generic EventHandler interface for EventManager.
// It is instantiated in event_manager.go but implemented here to keep files clean.
type KVEventHandler struct {
manager *EventManager
tenant_id string
// svcName string
modelName string
loraName string
instanceID string
blockSize int64
additionalSalt string
}
func (h *KVEventHandler) HandleEvent(event zmq.KVEvent, dpRank int64) error {
h.manager.mu.RLock()
if h.manager.stopped {
h.manager.mu.RUnlock()
return fmt.Errorf("manager stopped")
}
h.manager.mu.RUnlock()
slog.Info("Handling KV event", "instance_id", h.instanceID, "dpRank", dpRank)
// Create context for processing
ctx, cancel := context.WithTimeout(h.manager.ctx, 10*time.Second)
defer cancel()
// Dispatch event
switch e := event.(type) {
case *zmq.BlockStoredEvent:
slog.Debug("BlockStored",
"instance_id", h.instanceID,
"dpRank", dpRank,
"blocks", len(e.BlockHashes),
)
slog.Info("Received BlockStoredEvent", "medium", e.Medium)
return h.handleBlockStored(ctx, e, dpRank)
case *zmq.BlockRemovedEvent:
slog.Debug("BlockRemoved",
"instance_id", h.instanceID,
"dpRank", dpRank,
"blocks", len(e.BlockHashes),
)
slog.Info("Received BlockRemovedEvent", "medium", e.Medium)
return h.handleBlockRemoved(ctx, e, dpRank)
default:
slog.Warn("Unknown event type",
"type", fmt.Sprintf("%T", event),
)
return nil
}
}
func (h *KVEventHandler) handleBlockStored(ctx context.Context, event *zmq.BlockStoredEvent, dpRank int64) error {
// Convert to kvindexer event
conductorEvent := common.StoredEvent{
BlockHashes: event.BlockHashes,
BlockSize: event.BlockSize,
ModelName: h.modelName,
LoraName: h.loraName,
InstanceID: h.instanceID,
ParentBlockHash: event.ParentBlockHash,
TokenIds: event.TokenIDs,
Medium: event.Medium,
}
indexer := h.manager.getIndexer()
er := indexer.ProcessStoreEvent(conductorEvent, dpRank)
// TODO support mooncake_key map
if er != nil {
slog.Error("process store event failed.", "error", er)
}
slog.Debug("in handleBlockStored", "conductorEvent", conductorEvent)
return nil
}
func (h *KVEventHandler) handleBlockRemoved(ctx context.Context, event *zmq.BlockRemovedEvent, dpRank int64) error {
// Convert to conductor event
conductorEvent := common.RemovedEvent{
BlockHashes: event.BlockHashes,
ModelName: h.modelName,
LoraName: h.loraName,
InstanceID: h.instanceID,
BlockSize: h.blockSize,
Medium: event.Medium,
}
indexer := h.manager.getIndexer()
er := indexer.ProcessRemoveEvent(conductorEvent, dpRank, h.instanceID)
if er != nil {
slog.Error("process remove event failed.")
}
slog.Debug("in handleBlockRemoved", "conductorEvent", conductorEvent)
return nil
}
// TODO support mooncake update kv event

View File

@ -0,0 +1,492 @@
package kvevent
import (
"context"
"encoding/json"
"fmt"
"log/slog"
"net/http"
"sync"
"time"
"conductor/common"
"conductor/prefixindex"
"conductor/zmq"
)
// Dynamic register structure
type RegisterReq struct {
Endpoint string `json:"endpoint"`
ReplayEndpoint string `json:"replay_endpoint"`
Type string `json:"type"`
ModelName string `json:"modelname"`
LoraName *string `json:"lora_name"`
TenantID *string `json:"tenant_id"`
InstanceID string `json:"instance_id"`
BlockSize int64 `json:"block_size"`
DPRank int `json:"dp_rank"`
AdditionalSalt *string `json:"additionalsalt"`
}
// Dynamic unregister structure
type UnregisterReq struct {
Type string `json:"type"`
ModelName string `json:"modelname"`
LoraName *string `json:"lora_name"`
TenantID *string `json:"tenant_id"`
InstanceID string `json:"instance_id"`
BlockSize int `json:"block_size"`
DPRank int `json:"dp_rank"`
}
type QueryReq struct {
ModelName string `json:"model"`
LoraName *string `json:"lora_name"`
LoraID *int64 `json:"lora_id"`
TokenIDs []int32 `json:"token_ids"`
InstanceID *string `json:"instance_id"`
TenantID *string `json:"tenant_id"`
BlockSize int64 `json:"block_size"`
CacheSalt *string `json:"cache_salt"`
}
type EventManager struct {
indexer *prefixindex.PrefixCacheTable
services []common.ServiceConfig
httpserverport int
subscribers common.SyncMap[string, *zmq.ZMQClient]
// Map to store active configurations
activeConfigs common.SyncMap[string, common.ServiceConfig]
// Map to store tenant instance list
tenantInstanceMap map[string]map[string]struct{}
ctx context.Context
cancel context.CancelFunc
mu sync.RWMutex
tenantMutex sync.RWMutex
stopped bool
}
func NewEventManager(
services []common.ServiceConfig,
httpserverport int,
) *EventManager {
ctx, cancel := context.WithCancel(context.Background())
indexer := prefixindex.NewPrefixCacheTable()
// TODO 每个ModelContext创建一个独立的indexer
return &EventManager{
services: services,
indexer: indexer,
httpserverport: httpserverport,
ctx: ctx,
cancel: cancel,
tenantInstanceMap: make(map[string]map[string]struct{}),
}
}
func (m *EventManager) Start() error {
slog.Info("Starting KV Event Manager...")
// Subscribe to all services concurrently
var wg sync.WaitGroup
errCh := make(chan error, len(m.services))
for _, svc := range m.services {
wg.Add(1)
go func(service common.ServiceConfig) {
defer wg.Done()
if err := m.subscribeToService(service); err != nil {
slog.Error("Failed to initiate subscription",
"service_type", service.Type,
"instance_id", service.InstanceID,
"endpoint", service.Endpoint,
"error", err,
)
errCh <- fmt.Errorf("failed to subscribe to %s: %w", service.InstanceID, err)
}
}(svc)
}
wg.Wait()
close(errCh)
failureCount := len(errCh)
successCount := len(m.services) - failureCount
slog.Info("Static KV Event Manager started. Subscriptions",
"success", successCount,
"failed", failureCount,
)
return nil
}
func (m *EventManager) Stop() {
m.mu.Lock()
if m.stopped {
m.mu.Unlock()
return
}
m.stopped = true
m.mu.Unlock()
slog.Info("Stopping Conductor KV Event Manager.....")
// Cancel context
m.cancel()
// Stop all ZMQ clients
m.subscribers.Range(func(key string, client *zmq.ZMQClient) bool {
client.Stop()
slog.Info("Stopped all subscription",
"service_key", key,
)
return true
})
}
func makeServiceKey(instanceID string, tenantID string, dpRank int) string {
return fmt.Sprintf("%s|%s|%d", instanceID, tenantID, dpRank)
}
func (m *EventManager) subscribeToService(svc common.ServiceConfig) error {
// Use (instance_id, tenant_id) as composite key to support multi-tenant replicas
svcKey := makeServiceKey(svc.InstanceID, svc.TenantID, svc.DPRank)
if svc.InstanceID == "" {
svcKey = makeServiceKey(svc.Endpoint, svc.TenantID, svc.DPRank)
}
if _, exists := m.subscribers.Load(svcKey); exists {
return nil
}
// Validate endpoint
if svc.Endpoint == "" {
return fmt.Errorf("endpoint is required")
}
// Use ReplayEndpoint directly, fallback to empty if not provided
replayEndpoint := svc.ReplayEndpoint
handler := &KVEventHandler{
manager: m,
tenant_id: svc.TenantID,
modelName: svc.ModelName,
loraName: svc.LoraName,
instanceID: svc.InstanceID,
blockSize: svc.BlockSize,
additionalSalt: svc.AdditionalSalt,
}
// Configure ZMQ Client
zmqConfig := &zmq.ZMQClientConfig{
CachePoolKey: svcKey,
Endpoint: svc.Endpoint,
ReplayEndpoint: replayEndpoint,
ModelName: svc.ModelName,
PollTimeout: 100 * time.Millisecond,
ReplayTimeout: 5 * time.Second,
ReconnectDelay: 1 * time.Second,
}
if err := zmq.ValidateConfig(zmqConfig); err != nil {
return fmt.Errorf("invalid ZMQ config: %w", err)
}
client := zmq.NewZMQClient(zmqConfig, handler)
if err := client.Start(); err != nil {
return fmt.Errorf("failed to start ZMQ client: %w", err)
}
m.subscribers.Store(svcKey, client)
m.activeConfigs.Store(svcKey, svc)
//Add instance to tenant's instance map
m.tenantMutex.Lock()
if _, exists := m.tenantInstanceMap[svc.TenantID]; !exists {
m.tenantInstanceMap[svc.TenantID] = make(map[string]struct{})
}
m.tenantInstanceMap[svc.TenantID][svc.InstanceID] = struct{}{}
m.tenantMutex.Unlock()
slog.Info("Successfully subscribed to service",
"service_type", svc.Type,
"service_key", svcKey,
"instance_id", svc.InstanceID,
"tenant_id", svc.TenantID,
"endpoint", svc.Endpoint,
"replay_endpoint", replayEndpoint,
)
return nil
}
func (m *EventManager) unsubscribeFromService(instanceID string, tenantID string, dpRank int) {
svcKey := makeServiceKey(instanceID, tenantID, dpRank)
if client, exists := m.subscribers.Load(svcKey); exists {
client.Stop()
m.subscribers.Delete(svcKey)
m.activeConfigs.Delete(svcKey)
// Remove engine_instance from tenant's instance set
m.tenantMutex.Lock()
if instanceSet, exists := m.tenantInstanceMap[tenantID]; exists {
delete(instanceSet, instanceID)
}
m.tenantMutex.Unlock()
slog.Info("Successfully unsubscribed from service",
"service_key", svcKey,
"instance_id", instanceID,
"tenant_id", tenantID,
)
}
}
func (m *EventManager) getIndexer() *prefixindex.PrefixCacheTable {
return m.indexer
}
func (m *EventManager) StartHTTPServer() error {
mux := http.NewServeMux()
mux.HandleFunc("/query", func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
return
}
var req QueryReq
slog.Debug(
"receive req",
"method", r.Method,
"path", r.URL.Path,
"remote", r.RemoteAddr,
)
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
slog.Error("Failed to decode JSON", "err", err)
http.Error(w, "Invalid JSON", http.StatusBadRequest)
return
}
tenantID := "default"
if req.TenantID != nil && *req.TenantID != "" {
tenantID = *req.TenantID
}
loraName := ""
if req.LoraName != nil {
slog.Debug("LoraName is provided", "lora_name", *req.LoraName)
loraName = *req.LoraName
}
cacheSalt := ""
if req.CacheSalt != nil {
slog.Debug("cacheSalt is provided", "cacheSalt", *req.CacheSalt)
cacheSalt = *req.CacheSalt
}
response_result := make(map[string]map[string]prefixindex.CacheHitResult)
if req.InstanceID != nil {
slog.Info("search all engine instance for tenant. ", "instance_id", req.InstanceID)
modelContext := &prefixindex.ModelContext{
TenantID: tenantID,
ModelName: req.ModelName,
LoraName: loraName,
BlockSize: req.BlockSize,
AdditionalSalt: cacheSalt,
InstanceID: *req.InstanceID,
}
result := m.indexer.CacheHitCompute(modelContext, req.TokenIDs)
if result != nil {
if response_result[tenantID] == nil {
response_result[tenantID] = make(map[string]prefixindex.CacheHitResult)
}
response_result[tenantID][*req.InstanceID] = *result
}
} else {
if instanceSet, exists := m.tenantInstanceMap[tenantID]; exists {
for instanceID := range instanceSet {
modelContext := &prefixindex.ModelContext{
TenantID: tenantID,
ModelName: req.ModelName,
LoraName: loraName,
BlockSize: req.BlockSize,
AdditionalSalt: cacheSalt,
InstanceID: instanceID,
}
result := m.indexer.CacheHitCompute(modelContext, req.TokenIDs)
if result != nil {
if response_result[tenantID] == nil {
response_result[tenantID] = make(map[string]prefixindex.CacheHitResult)
}
response_result[tenantID][instanceID] = *result
}
}
} else {
slog.Warn("current tenant has no engine_instance. ", "tenant_id", tenantID)
}
}
slog.Debug("cache hit status", "hitresult", response_result)
w.Header().Set("Content-Type", "application/json")
if err := json.NewEncoder(w).Encode(response_result); err != nil {
slog.Error("Failed to encode response", "err", err)
http.Error(w, "Internal Server Error", http.StatusInternalServerError)
return
}
})
// Register interface
mux.HandleFunc("/register", func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
return
}
var req RegisterReq
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
slog.Error("Failed to decode register JSON", "err", err)
http.Error(w, "Invalid JSON", http.StatusBadRequest)
return
}
// Handle Optional fields' default values
tenantID := "default"
if req.TenantID != nil && *req.TenantID != "" {
tenantID = *req.TenantID
}
loraName := ""
if req.LoraName != nil {
slog.Info("LoraName is provided", "lora_name", *req.LoraName)
loraName = *req.LoraName
}
additionalSalt := ""
if req.AdditionalSalt != nil {
additionalSalt = *req.AdditionalSalt
}
svc := common.ServiceConfig{
Endpoint: req.Endpoint,
ReplayEndpoint: req.ReplayEndpoint,
Type: req.Type,
ModelName: req.ModelName,
LoraName: loraName,
TenantID: tenantID,
InstanceID: req.InstanceID,
BlockSize: req.BlockSize,
DPRank: req.DPRank,
AdditionalSalt: additionalSalt,
}
if err := m.subscribeToService(svc); err != nil {
slog.Error("Dynamic register failed", "instance_id", req.InstanceID, "err", err)
http.Error(w, fmt.Sprintf("Failed to subscribe: %v", err), http.StatusInternalServerError)
return
}
m.services = append(m.services, svc)
modelContext := &prefixindex.ModelContext{
TenantID: tenantID,
ModelName: req.ModelName,
LoraName: loraName,
BlockSize: req.BlockSize,
AdditionalSalt: additionalSalt,
InstanceID: svc.InstanceID,
}
m.indexer.AddDpSize(modelContext, int64(svc.DPRank))
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(map[string]string{
"status": "registered successfully",
"instance_id": svc.InstanceID,
})
})
// Unregister interface
mux.HandleFunc("/unregister", func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
return
}
var req UnregisterReq
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
slog.Error("Failed to decode unregister JSON", "err", err)
http.Error(w, "Invalid JSON", http.StatusBadRequest)
return
}
// Build target service key from instance_id and tenant_id
targetTenant := "default"
if req.TenantID != nil && *req.TenantID != "" {
targetTenant = *req.TenantID
}
targetKey := makeServiceKey(req.InstanceID, targetTenant, req.DPRank)
// Direct lookup and removal
if _, exists := m.activeConfigs.Load(targetKey); !exists {
http.Error(w, fmt.Sprintf("service not found: %s", targetKey), http.StatusNotFound)
return
}
m.unsubscribeFromService(req.InstanceID, targetTenant, req.DPRank)
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(map[string]interface{}{
"status": "unregistered successfully",
"removed_instances": []string{targetKey},
})
})
// Global view interface
mux.HandleFunc("/global_view", func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
return
}
globalView := m.indexer.GetGlobalView()
w.Header().Set("Content-Type", "application/json")
if err := json.NewEncoder(w).Encode(globalView); err != nil {
slog.Error("Failed to encode global view response", "err", err)
http.Error(w, "Internal Server Error", http.StatusInternalServerError)
return
}
})
server := &http.Server{
Addr: fmt.Sprintf(":%d", m.httpserverport),
Handler: mux,
}
go func() {
slog.Info("HTTP server listening", "port", m.httpserverport)
if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
slog.Error("HTTP server failed", "err", err)
}
}()
// Start a goroutine to listen for context cancellation, used for graceful shutdown.
go func() {
<-m.ctx.Done()
slog.Info("Shutting down HTTP server")
// 5-second timeout for forced shutdown
shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if err := server.Shutdown(shutdownCtx); err != nil {
slog.Error("HTTP server shutdown error", "err", err)
server.Close()
}
}()
return nil
}

View File

@ -0,0 +1,131 @@
package main
import (
"encoding/json"
"errors"
"log/slog"
"os"
"os/signal"
"syscall"
"conductor/common"
"conductor/kvevent"
)
var (
// TODO change default config path
conductorConfigPath = common.LoadEnv("CONDUCTOR_CONFIG_PATH", "/root/conductor_config.json")
httpServerPort = 13333
)
type configStruct struct {
KVEventInstance map[string]serviceRaw `json:"kvevent_instance"`
HTTPPort int `json:"http_server_port"`
}
type serviceRaw struct {
Endpoint string `json:"endpoint"`
ReplayEndpoint string `json:"replay_endpoint"`
TypeStr string `json:"type"`
ModelName string `json:"modelname"`
LoraName string `json:"lora_name"`
TenantID string `json:"tenant_id"`
InstanceID string `json:"instance_id"`
BlockSize int64 `json:"block_size"`
DPRank int `json:"dp_rank"`
AdditionalSalt string `json:"additionalsalt"`
}
func mapServiceType(s string) (string, bool) {
switch s {
case "vLLM":
return common.ServiceTypeVLLM, true
case "Mooncake":
return common.ServiceTypeMooncake, true
default:
return "None", false
}
}
func parseConfig() []common.ServiceConfig {
if _, err := os.Stat(conductorConfigPath); errors.Is(err, os.ErrNotExist) {
slog.Warn("Config file does not exist, exiting.", "path", conductorConfigPath)
// os.Exit(1)
return []common.ServiceConfig{}
} else if err != nil {
slog.Warn("Error accessing config file", "path", conductorConfigPath, "error", err)
// os.Exit(1)
return []common.ServiceConfig{}
}
data, err := os.ReadFile(conductorConfigPath)
if err != nil {
slog.Error("Failed to read config file", "path", conductorConfigPath, "error", err)
os.Exit(1)
}
var cfg configStruct
if err := json.Unmarshal(data, &cfg); err != nil {
slog.Error("Failed to parse JSON config", "error", err)
os.Exit(1)
}
httpServerPort = cfg.HTTPPort
services := make([]common.ServiceConfig, 0, len(cfg.KVEventInstance))
for name, raw := range cfg.KVEventInstance {
serviceType, ok := mapServiceType(raw.TypeStr)
if !ok {
slog.Error("Unknown service type", "type", raw.TypeStr)
continue
}
services = append(services, common.ServiceConfig{
Endpoint: raw.Endpoint,
ReplayEndpoint: raw.ReplayEndpoint,
Type: serviceType,
ModelName: raw.ModelName,
LoraName: raw.LoraName,
TenantID: raw.TenantID,
InstanceID: name,
BlockSize: raw.BlockSize,
DPRank: raw.DPRank,
AdditionalSalt: raw.AdditionalSalt,
})
}
return services
}
func main() {
// TODO support print metrics for conductor
logLevel := common.ParseLogLevel()
logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{
Level: logLevel,
}))
slog.SetDefault(logger)
slog.Info("Starting Conductor KV Event Manager...", "logLevel", logLevel)
services := parseConfig()
manager := kvevent.NewEventManager(services, httpServerPort)
if err := manager.StartHTTPServer(); err != nil {
slog.Error("Failed to start HTTP server", "err", err)
}
if err := manager.Start(); err != nil {
slog.Error("Failed to start manager", "error", err)
os.Exit(1)
}
sigChan := make(chan os.Signal, 1)
signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
slog.Info("Manager is running. Press Ctrl+C to stop.")
<-sigChan
slog.Info("Shutting down...")
manager.Stop()
}

View File

@ -0,0 +1,561 @@
package prefixindex
import (
"encoding/binary"
"fmt"
"log/slog"
"math/rand"
"sync"
"sync/atomic"
"time"
"conductor/common"
"github.com/cespare/xxhash/v2"
)
var (
enableCpuEviction = common.LoadBoolEnv("ENABLE_CPU_EVICTION", false)
maxCpuKeyNum = int64(common.LoadIntEnv("MAX_CPU_KEY_NUM", 30000))
cpuKeyEvictionRatio = common.LoadFloatEnv("CPU_KEY_EVICTION_RATIO", 0.2)
// TODO
// The eviction parameter here is used to control that the number of CPU cache prefixes cannot grow without limit.
// The best way is to add an zmq publisher in mooncake-store to actively notify conductor of the block eviction.
)
type ModelContext struct {
ModelName string
LoraName string // None represents no LoRA adapter
BlockSize int64
// TODO @yejj710
// Confirm the difference between the previously discussed cache_salt and additionalSalt.
//The current understanding is that cache_salt is used to ensure data isolation between different customers,
// and it seems it can be directly added to additionalSalt.
AdditionalSalt string
TenantID string
InstanceID string // unique identifier for each API server
}
type CacheStoreInfo struct {
// TODO Currently, the KV cache at different levels is not distinguished.
// In the future, the caches of Mooncake and inference engines (vLLM, SGLang)
// should be handled separately.
engineLastAccessTime atomic.Int64
TotalReplicaNums atomic.Int64
mediumSet map[string]struct{}
dpRankSet map[int64]struct{} // indicate the dp_rank that the block is cached on
// LRU linked list pointers
lruPrev *CacheStoreInfo
lruNext *CacheStoreInfo
}
type HashMapStore struct {
// conductor prefixHash -> cachestore
prefixMap map[uint64]*CacheStoreInfo
createTime time.Time
lastAccess atomic.Int64
totalPrefixes int64
// LRU linked list: head is least recently used, tail is most recently used
lruHead *CacheStoreInfo
lruTail *CacheStoreInfo
}
type ContextData struct {
prefixMu sync.RWMutex
hashmapMu sync.RWMutex
prefixStore *HashMapStore
seed uint64
DpSize map[int64]struct{}
proxyHashMapping map[uint64]uint64 // engine block hash -> conductor prefix hash
}
type PrefixCacheTable struct {
// TODO use instance_id to distinguish different engine instances
contextMap sync.Map // ModelContext → *ContextData
contextCount atomic.Int32
}
type CacheHitResult struct {
LongestMatchTokens int64 `json:"longest_matched"`
DP map[int64]int64 `json:"DP"`
GPU int64 `json:"GPU"`
CPU int64 `json:"CPU"`
DISK int64 `json:"DISK"`
}
type ModelContextView struct {
ModelName string `json:"model_name"`
LoraName string `json:"lora_name"`
BlockSize int64 `json:"block_size"`
AdditionalSalt string `json:"additional_salt"`
TenantID string `json:"tenant_id"`
InstanceID string `json:"instance_id"`
}
type GlobalView struct {
ContextCount int32 `json:"context_count"`
ModelContexts []ModelContextView `json:"model_contexts"`
ProxyHashMap []map[uint64]uint64 `json:"hashmap"`
}
func GenerateSeedFromEnv() uint64 {
r := rand.New(rand.NewSource(time.Now().Unix()))
envSeed := common.LoadIntEnv("CONDUCTOR_SEED", -1)
var seed uint64
if envSeed != -1 {
seed = uint64(envSeed)
} else {
seed = r.Uint64()
}
return seed
}
func NewPrefixCacheTable() *PrefixCacheTable {
p := &PrefixCacheTable{}
return p
}
func (p *PrefixCacheTable) getContextData(modelcontext *ModelContext) *ContextData {
ctx_value := *modelcontext
value, exists := p.contextMap.Load(ctx_value)
if exists {
return value.(*ContextData)
}
seedValue := xxhash.Sum64String(modelcontext.AdditionalSalt)
newContextData := &ContextData{
prefixStore: &HashMapStore{
prefixMap: make(map[uint64]*CacheStoreInfo),
createTime: time.Now(),
totalPrefixes: 0,
},
proxyHashMapping: make(map[uint64]uint64),
seed: seedValue,
DpSize: make(map[int64]struct{}),
}
newContextData.prefixStore.lastAccess.Store(time.Now().Unix())
p.contextMap.Store(ctx_value, newContextData)
slog.Debug("in getContextData", "modelcontext", modelcontext)
p.contextCount.Add(1)
return newContextData
}
func (p *PrefixCacheTable) AddDpSize(modelcontext *ModelContext, dpRank int64) {
// value, exists := p.contextMap.Load(modelcontext)
contextData := p.getContextData(modelcontext)
contextData.DpSize[dpRank] = struct{}{}
}
func (p *PrefixCacheTable) ComputePrefixHash(modelcontext *ModelContext, tokenIds []int32, cacheSalt uint64) []uint64 {
// cacheSalt is used to seperate hash from different customers
numBlocks := len(tokenIds) / int(modelcontext.BlockSize)
prefixHashes := make([]uint64, 0, numBlocks)
var parentHash uint64 = cacheSalt
for i := 0; i < numBlocks; i++ {
start := i * int(modelcontext.BlockSize)
end := start + int(modelcontext.BlockSize)
if end > len(tokenIds) {
break
}
hashValue := p.computeHash(parentHash, tokenIds[start:end])
prefixHashes = append(prefixHashes, hashValue)
parentHash = hashValue
}
return prefixHashes
}
func (p *PrefixCacheTable) CacheHitCompute(modelcontext *ModelContext, tokenIds []int32) *CacheHitResult {
value, exists := p.contextMap.Load(*modelcontext)
prefixMatchResult := &CacheHitResult{
LongestMatchTokens: 0,
DP: map[int64]int64{},
GPU: 0,
CPU: 0,
DISK: 0,
}
if !exists {
slog.Error("In CacheHitCompute, contextData not found")
return prefixMatchResult
}
contextData := value.(*ContextData)
cacheSalt := xxhash.Sum64String(modelcontext.AdditionalSalt)
prefixHashes := p.ComputePrefixHash(modelcontext, tokenIds, cacheSalt)
// TODO @yejj710
// When there is no data in contextData, what information should be returned for the matched modelcontext
// This is related to function `AddDpSize`
slog.Debug("In CacheHitCompute", "prefixHashes", prefixHashes)
contextData.prefixMu.RLock()
defer contextData.prefixMu.RUnlock()
prefixStore := contextData.prefixStore
// reserve prefixHashes and then compute cache hit
for _, prefixHash := range prefixHashes {
cacheStoreInfo, exists := prefixStore.prefixMap[prefixHash]
slog.Debug("In CacheHitCompute", "cacheStoreInfo", cacheStoreInfo)
// chained hash, break if no replica exists
if !exists || cacheStoreInfo.TotalReplicaNums.Load() == 0 {
break
}
cacheHit := false
for key := range cacheStoreInfo.mediumSet {
if key == "cpu" {
prefixMatchResult.CPU += modelcontext.BlockSize
cacheHit = true
} else if key == "GPU" {
prefixMatchResult.GPU += modelcontext.BlockSize
cacheHit = true
} else {
slog.Warn("In CacheHitCompute, unknown medium type", "medium", key)
}
}
if cacheHit {
prefixMatchResult.LongestMatchTokens += modelcontext.BlockSize
for dpRank := range cacheStoreInfo.dpRankSet {
prefixMatchResult.DP[dpRank] += modelcontext.BlockSize
}
cacheStoreInfo.engineLastAccessTime.Store(time.Now().Unix())
// move to tail (most recently used)
if enableCpuEviction {
prefixStore.addToLRUTail(cacheStoreInfo)
// TODO update LRU list asynchronously
}
}
}
prefixStore.lastAccess.Store(time.Now().Unix())
return prefixMatchResult
}
func (p *PrefixCacheTable) ProcessStoreEvent(event common.StoredEvent, dpRank int64) error {
if len(event.BlockHashes) == 0 {
return nil
}
tenantID := "default"
slog.Debug("In ProcessStoreEvent", "modelName", event.ModelName, "instanceID", event.InstanceID, "dpRank", dpRank)
contextData := p.getContextData(&ModelContext{
ModelName: event.ModelName,
LoraName: event.LoraName,
BlockSize: event.BlockSize,
TenantID: tenantID,
AdditionalSalt: "",
InstanceID: event.InstanceID,
})
contextData.hashmapMu.Lock()
defer contextData.hashmapMu.Unlock()
proxyHashMap := contextData.proxyHashMapping
if len(event.BlockHashes)*int(event.BlockSize) != len(event.TokenIds) {
if len(event.BlockHashes) != 1 {
return fmt.Errorf("block hashes and tokens length mismatch")
}
}
newPrefixStore := make([]struct {
hashValue uint64
engineID string
}, 0)
var parentHash uint64 = contextData.seed
// TODO If the ParentBlockHash happens to be 0, a bug will occur here, because 0 is a valid hash value.
if event.ParentBlockHash != 0 {
slog.Debug("parent Block HASH is not None.")
if pbh, exists := proxyHashMap[event.ParentBlockHash]; exists {
parentHash = pbh
}
}
for i, blockHash := range event.BlockHashes {
// cache already exists, add engine info and continue
if existingHash, exists := proxyHashMap[blockHash]; exists {
newPrefixStore = append(newPrefixStore, struct {
hashValue uint64
engineID string
}{existingHash, event.InstanceID})
continue
}
// if not exists, compute hash
hashValue := p.computeHash(parentHash, event.TokenIds[i*int(event.BlockSize):(i+1)*int(event.BlockSize)])
parentHash = hashValue
proxyHashMap[blockHash] = hashValue
newPrefixStore = append(newPrefixStore, struct {
hashValue uint64
engineID string
}{
hashValue: hashValue,
engineID: event.InstanceID,
})
}
if len(newPrefixStore) > 0 {
contextData.prefixMu.Lock()
defer contextData.prefixMu.Unlock()
prefixStore := contextData.prefixStore
for _, newPrefix := range newPrefixStore {
slog.Debug("show new prefix data", "newPrefix", newPrefix)
p.addNewPrefixStore(prefixStore, newPrefix.hashValue, newPrefix.engineID, event.Medium, dpRank)
}
if enableCpuEviction && prefixStore.totalPrefixes > maxCpuKeyNum {
evictedCount := prefixStore.evictLRU(cpuKeyEvictionRatio)
slog.Info("LRU eviction triggered", "totalPrefixes", prefixStore.totalPrefixes, "evictedCount", evictedCount)
}
}
return nil
}
func (p *PrefixCacheTable) ProcessRemoveEvent(event common.RemovedEvent, dpRank int64, instanceID string) error {
if len(event.BlockHashes) == 0 {
return nil
}
contextData := p.getContextData(&ModelContext{
ModelName: event.ModelName,
LoraName: event.LoraName,
BlockSize: event.BlockSize,
TenantID: "default",
AdditionalSalt: "",
InstanceID: instanceID,
})
contextData.hashmapMu.Lock()
defer contextData.hashmapMu.Unlock()
proxyHashMap := contextData.proxyHashMapping
removeConductorHash := make([]uint64, 0, len(event.BlockHashes))
// delete proxyHashMapping
contextData.prefixMu.Lock()
defer contextData.prefixMu.Unlock()
prefixStore := contextData.prefixStore
for _, blockHash := range event.BlockHashes {
if conductorHash, exists := proxyHashMap[blockHash]; exists {
removeConductorHash = append(removeConductorHash, conductorHash)
// Only delete proxyHashMapping entry when all replicas are removed
cacheStoreInfo, exists := prefixStore.prefixMap[conductorHash]
if !exists {
continue
}
if cacheStoreInfo.TotalReplicaNums.Load() == 1 {
delete(proxyHashMap, blockHash)
}
}
}
// update prefixStore
for _, conductorHash := range removeConductorHash {
cacheStoreInfo, exists := prefixStore.prefixMap[conductorHash]
if !exists {
continue
}
cacheStoreInfo.TotalReplicaNums.Add(-1)
// Remove per-instance metadata
delete(cacheStoreInfo.dpRankSet, dpRank)
delete(cacheStoreInfo.mediumSet, event.Medium)
slog.Info("process remove event", "conductorHash", conductorHash, "dpRank", dpRank, "medium", event.Medium)
// Only delete entry when all replicas are removed
if cacheStoreInfo.TotalReplicaNums.Load() <= 0 {
delete(prefixStore.prefixMap, conductorHash)
prefixStore.totalPrefixes--
}
}
return nil
}
func (p *PrefixCacheTable) computeHash(parentHash uint64, blockTokenIDs []int32) uint64 {
// digest := xxhash.NewWithSeed()
digest := xxhash.New()
var parentHashBytes [8]byte
binary.LittleEndian.PutUint64(parentHashBytes[:], parentHash)
_, _ = digest.Write(parentHashBytes[:])
var tokenIDsBytes [8]byte
for _, tokenID := range blockTokenIDs {
binary.LittleEndian.PutUint32(tokenIDsBytes[:], uint32(tokenID))
_, _ = digest.Write(tokenIDsBytes[:])
}
return digest.Sum64()
}
func (p *PrefixCacheTable) addNewPrefixStore(prefixStore *HashMapStore, hashValue uint64, instanceID string, medium string, dpRank int64) {
now := time.Now().Unix()
if prefixStore.prefixMap[hashValue] == nil {
slog.Debug("in addNewPrefixStore, prefixStore.prefixMap[hashValue] is nil", "hashValue", hashValue)
prefixStore.prefixMap[hashValue] = &CacheStoreInfo{
mediumSet: make(map[string]struct{}),
dpRankSet: make(map[int64]struct{}),
}
prefixStore.totalPrefixes++
}
cacheStoreInfo := prefixStore.prefixMap[hashValue]
cacheStoreInfo.engineLastAccessTime.Store(now)
cacheStoreInfo.TotalReplicaNums.Add(1)
// TODO If using Mooncake-Store, you do not need to set dpRank, because it does not distinguish between kv-blocks of different dpRanks.
cacheStoreInfo.mediumSet[medium] = struct{}{}
cacheStoreInfo.dpRankSet[dpRank] = struct{}{}
slog.Debug("in addNewPrefixStore", "conductor_hash", hashValue, "current_mediumset", cacheStoreInfo.mediumSet[medium])
// move to tail (most recently used)
if enableCpuEviction {
prefixStore.addToLRUTail(cacheStoreInfo)
}
}
func (p *PrefixCacheTable) GetGlobalView() *GlobalView {
view := &GlobalView{
ContextCount: p.contextCount.Load(),
ModelContexts: make([]ModelContextView, 0),
ProxyHashMap: make([]map[uint64]uint64, 0),
}
p.contextMap.Range(func(key, value interface{}) bool {
ctx := key.(ModelContext)
contextData := value.(*ContextData)
ctxView := ModelContextView{
ModelName: ctx.ModelName,
LoraName: ctx.LoraName,
BlockSize: ctx.BlockSize,
AdditionalSalt: ctx.AdditionalSalt,
TenantID: ctx.TenantID,
InstanceID: ctx.InstanceID,
}
contextData.prefixMu.RLock()
defer contextData.prefixMu.RUnlock()
contextData.hashmapMu.RLock()
defer contextData.hashmapMu.RUnlock()
view.ProxyHashMap = append(view.ProxyHashMap, contextData.proxyHashMapping)
view.ModelContexts = append(view.ModelContexts, ctxView)
return true
})
return view
}
// move or add a CacheStoreInfo to the tail of the LRU list (most recently used)
func (h *HashMapStore) addToLRUTail(cacheNode *CacheStoreInfo) {
if cacheNode == nil {
return
}
// If already in list, remove it first
if cacheNode.lruPrev != nil || cacheNode.lruNext != nil || h.lruHead == cacheNode {
h.removeFromLRU(cacheNode)
}
if h.lruTail == nil {
// Empty list
h.lruHead = cacheNode
h.lruTail = cacheNode
cacheNode.lruPrev = nil
cacheNode.lruNext = nil
} else {
h.lruTail.lruNext = cacheNode
cacheNode.lruPrev = h.lruTail
cacheNode.lruNext = nil
h.lruTail = cacheNode
}
}
// remove a CacheStoreInfo from the LRU list
func (h *HashMapStore) removeFromLRU(cacheNode *CacheStoreInfo) {
if cacheNode == nil {
return
}
if cacheNode.lruPrev != nil {
cacheNode.lruPrev.lruNext = cacheNode.lruNext
} else {
// cacheNode is head
h.lruHead = cacheNode.lruNext
}
if cacheNode.lruNext != nil {
cacheNode.lruNext.lruPrev = cacheNode.lruPrev
} else {
// cacheNode is tail
h.lruTail = cacheNode.lruPrev
}
cacheNode.lruPrev = nil
cacheNode.lruNext = nil
}
// evict the least recently used cpu mediums based on eviction ratio
func (h *HashMapStore) evictLRU(ratio float64) int {
if h.lruHead == nil || ratio <= 0 {
return 0
}
// TODO Here it is assumed that each CacheStoreInfo contains a cpu medium, but in practice such assumptions should not be made.
// Instead, a separate variable should be used to record the number of cpu mediums.
targetCount := int(float64(h.totalPrefixes) * ratio)
evictedCount := 0
// Traverse from head (least recently used)
for h.lruHead != nil && evictedCount < targetCount {
cacheNode := h.lruHead
// Only evict `cpu` medium
if _, exist := cacheNode.mediumSet["cpu"]; exist {
delete(cacheNode.mediumSet, "cpu")
// If no mediums left, remove the entry entirely
if len(cacheNode.mediumSet) == 0 {
h.removeFromLRU(cacheNode)
h.deleteCacheStoreInfo(cacheNode)
// TODO remove proxyHashMapping in ContextData,
// currently we do not maintain reverse mapping from conductor hash to proxy hash,
// which makes it hard to delete the proxyHashMapping entry.
h.totalPrefixes--
}
} else {
// Move to LRU-List tail
h.removeFromLRU(cacheNode)
h.addToLRUTail(cacheNode)
}
evictedCount++
}
return evictedCount
}
func (h *HashMapStore) deleteCacheStoreInfo(cacheNode *CacheStoreInfo) {
for k, v := range h.prefixMap {
if v == cacheNode {
delete(h.prefixMap, k)
return
}
}
}

View File

@ -0,0 +1,106 @@
package zmq
import "time"
type EventType string
const (
EventTypeBlockStored EventType = "BlockStored"
EventTypeBlockRemoved EventType = "BlockRemoved"
// EventTypeBlockUpdate indicates that blocks have been updated from the KV cache
EventTypeBlockUpdate EventType = "BlockUpdate"
EventTypeAllCleared EventType = "AllBlocksCleared"
)
type KVEvent interface {
GetType() EventType
GetTimestamp() time.Time
}
type BlockStoredEvent struct {
Type EventType
Timestamp time.Time
BlockHashes []uint64
TokenIDs []int32
ParentBlockHash uint64
BlockSize int64
MooncakeKey string
ReplicaList [][]string
ModelName string
LoraID int64
LoraName string
PodName string
Medium string
}
func (e *BlockStoredEvent) GetType() EventType {
return e.Type
}
func (e *BlockStoredEvent) GetTimestamp() time.Time {
return e.Timestamp
}
type BlockRemovedEvent struct {
Type EventType
Timestamp time.Time
BlockHashes []uint64
ModelName string
PodName string
Medium string
}
func (e *BlockRemovedEvent) GetType() EventType {
return e.Type
}
func (e *BlockRemovedEvent) GetTimestamp() time.Time {
return e.Timestamp
}
type AllBlocksClearedEvent struct {
Type EventType
Timestamp time.Time
ModelName string
PodName string
}
func (e *AllBlocksClearedEvent) GetType() EventType {
return e.Type
}
func (e *AllBlocksClearedEvent) GetTimestamp() time.Time {
return e.Timestamp
}
type BlockUpdateEvent struct {
Type EventType
Timestamp time.Time
BlockHashes []uint64
TokenIDs []int32
ParentBlockHash uint64
ModelName string
PodName string
BlockSize int64
}
// GetType returns the event type
func (e *BlockUpdateEvent) GetType() EventType {
return e.Type
}
func (e *BlockUpdateEvent) GetTimestamp() time.Time {
return e.Timestamp
}
const (
SourceMooncake string = "mooncake"
SourceVLLM string = "vllm"
)
type EventBatch struct {
Source string // indicates the origin of the event batch
Events []KVEvent
DataParallelRank int64
}

View File

@ -0,0 +1,676 @@
package zmq
import (
"fmt"
"log/slog"
"strconv"
"strings"
"time"
msgpack "github.com/shamaton/msgpack/v2"
)
type EventParser interface {
ParseEvent(raw []interface{}, timestamp interface{}) (KVEvent, error)
EventMappings() map[string]EventType
Source() string
}
type mooncakeParser struct{}
func (p *mooncakeParser) Source() string { return SourceMooncake }
func (p *mooncakeParser) EventMappings() map[string]EventType {
return map[string]EventType{
"BlockStoreEvent": EventTypeBlockStored,
"BlockUpdateEvent": EventTypeBlockUpdate,
"RemoveAllEvent": EventTypeAllCleared,
}
}
func (p *mooncakeParser) ParseEvent(raw []interface{}, timestamp interface{}) (KVEvent, error) {
eventTypeStr, ok := raw[0].(string)
if !ok {
return nil, fmt.Errorf("invalid event type format: %T", raw[0])
}
eventType, exists := p.EventMappings()[eventTypeStr]
if !exists {
return nil, fmt.Errorf("unknown mooncake event type: %s", eventTypeStr)
}
switch eventType {
case EventTypeBlockStored:
return parseMooncakeBlockStored(raw, timestamp)
default:
return nil, fmt.Errorf("unhandled event: %s", eventType)
}
}
func decodeCommonEventBatch(
data []byte,
expectedLength int,
extractEvents func([]interface{}) ([]interface{}, interface{}, error),
parser EventParser,
) (*EventBatch, error) {
if len(data) > 0 {
slog.Debug("First byte of payload", "hex", fmt.Sprintf("%02x", data[0]))
}
var arr []interface{}
if err := msgpack.Unmarshal(data, &arr); err != nil {
return nil, fmt.Errorf("failed to unmarshal event batch: %w", err)
}
if len(arr) != expectedLength {
return nil, fmt.Errorf("expected %d-element array, got %d", expectedLength, len(arr))
}
events, timestamp, err := extractEvents(arr)
if err != nil {
return nil, err
}
if len(events) == 0 {
slog.Warn("Received empty event list")
}
var dpRank int64 = -1
if expectedLength == 3 {
dpRank, err = parseInt64(arr[2])
if err != nil {
return nil, fmt.Errorf("failed to parse dpRank: %w", err)
}
}
batch := &EventBatch{
Source: parser.Source(),
Events: make([]KVEvent, 0, len(events)),
DataParallelRank: dpRank,
}
slog.Info("Receive batched kv-event", "source", batch.Source, "dpRank", dpRank)
for i, rawEvent := range events {
eventSlice, ok := rawEvent.([]interface{})
if !ok {
return nil, fmt.Errorf("event at index %d is not a slice: %T", i, rawEvent)
}
event, err := parser.ParseEvent(eventSlice, timestamp)
if err != nil {
return nil, fmt.Errorf("failed to parse event at index %d: %w", i, err)
}
batch.Events = append(batch.Events, event)
}
return batch, nil
}
func newMooncakeParser() EventParser {
return &mooncakeParser{}
}
func DecodeMooncakeEventBatch(data []byte) (*EventBatch, error) {
return decodeCommonEventBatch(
data,
2,
func(arr []interface{}) ([]interface{}, interface{}, error) {
events, ok := arr[1].([]interface{})
if !ok {
return nil, nil, fmt.Errorf("invalid events type: %T", arr[1])
}
return events, arr[0], nil
},
newMooncakeParser(),
)
}
type vllmParser struct{}
func (p *vllmParser) Source() string { return SourceVLLM }
func (p *vllmParser) EventMappings() map[string]EventType {
return map[string]EventType{
"BlockStored": EventTypeBlockStored,
"BlockRemoved": EventTypeBlockRemoved,
"AllBlocksCleared": EventTypeAllCleared,
}
}
func (p *vllmParser) ParseEvent(raw []interface{}, timestamp interface{}) (KVEvent, error) {
eventTypeStr, ok := raw[0].(string)
if !ok {
return nil, fmt.Errorf("invalid event type format: %T", raw[0])
}
eventType, exists := p.EventMappings()[eventTypeStr]
if !exists {
return nil, fmt.Errorf("unknown vllm event type: %s", eventTypeStr)
}
switch eventType {
case EventTypeBlockStored:
return parseVllmBlockStored(raw, timestamp)
case EventTypeBlockRemoved:
return parseVllmBlockRemoved(raw, timestamp)
default:
return nil, fmt.Errorf("unhandled event: %s", eventType)
}
}
func newVLLMParser() EventParser {
return &vllmParser{}
}
func DecodeVllmEventBatch(data []byte) (*EventBatch, error) {
return decodeCommonEventBatch(
data,
3,
func(arr []interface{}) ([]interface{}, interface{}, error) {
events, ok := arr[1].([]interface{})
if !ok {
return nil, nil, fmt.Errorf("invalid events type: %T", arr[1])
}
return events, arr[0], nil
},
newVLLMParser(),
)
}
func parseMooncakeBlockStored(data []interface{}, timestamp interface{}) (*BlockStoredEvent, error) {
event := &BlockStoredEvent{
Type: EventTypeBlockStored,
}
if mooncakekey, err := safeGetString(data[1]); err == nil {
event.MooncakeKey = mooncakekey
} else {
return nil, fmt.Errorf("failed to parse MooncakeKey from field at index 1: %w", err)
}
if replicalist, err := convertToReplicaList(data[2]); err == nil {
event.ReplicaList = replicalist
slog.Debug("ReplicaList:", "ReplicaList", event.ReplicaList)
} else {
return nil, fmt.Errorf("failed to parse ReplicaList from field at index 2: %w", err)
}
if blocksize, err := parseInt64(data[4]); err == nil {
event.BlockSize = blocksize
slog.Debug("BlockSize:", "BlockSize", event.BlockSize)
} else {
return nil, fmt.Errorf("failed to parse BlockSize from field at index 4: %w", err)
}
if blockhash, err := parseMooncakeParentUint64(data[5]); err == nil {
event.BlockHashes = blockhash
slog.Debug("BlockHashes:", "BlockHashes", event.BlockHashes)
} else {
return nil, fmt.Errorf("failed to parse BlockHashes from field at index 5: %w", err)
}
if parentblockhash, err := parseMooncakeUint64(data[6]); err == nil {
event.ParentBlockHash = parentblockhash
slog.Debug("ParentBlockHash:", "ParentBlockHash", event.ParentBlockHash)
} else {
return nil, fmt.Errorf("failed to parse ParentBlockHash from field at index 6: %w", err)
}
if tokenIDsRaw, ok := data[7].([]interface{}); ok {
tokens, err := parseInt32Array(tokenIDsRaw)
if err != nil {
return nil, fmt.Errorf("failed to parse TokenIDs from field at index 7: %w", err)
}
event.TokenIDs = tokens
slog.Debug("TokenIDs:", "TokenIDs", event.TokenIDs)
} else {
return nil, fmt.Errorf("missing or invalid token_ids")
}
return event, nil
}
// parseBlockStoredEvent parses a BlockStoredEvent from raw data
func parseVllmBlockStored(data []interface{}, timestamp interface{}) (*BlockStoredEvent, error) {
event := &BlockStoredEvent{
Type: EventTypeBlockStored,
}
for i, elem := range data {
slog.Debug("in parseVllmBlockStored:", "index", i, "type", fmt.Sprintf("%T", elem), "value", elem)
}
slog.Debug("in parseVllmBlockStored:", "timestamp", timestamp)
// Parse timestamp
if ts, err := parseTimestamp(timestamp); err == nil {
event.Timestamp = ts
} else {
return nil, fmt.Errorf("failed to parse timestamp: %w", err)
}
// Parse block hashes
if hashes, err := parseUint64Array(data[1]); err == nil {
event.BlockHashes = hashes
} else {
return nil, fmt.Errorf("failed to parse block_hashes: %w", err)
}
// Parse token IDs (array of arrays)
if tokenIDsRaw, ok := data[3].([]interface{}); ok {
tokens, err := parseInt32Array(tokenIDsRaw)
if err != nil {
return nil, fmt.Errorf("failed to parse token_ids at index %w", err)
}
event.TokenIDs = tokens
} else {
return nil, fmt.Errorf("missing or invalid token_ids")
}
var parentHash uint64
if data[2] == nil {
parentHash = uint64(0)
} else {
hash := data[2]
// fmt.Printf("Type of hash>>>>: %T\n", hash)
if h, ok := hash.(uint64); ok {
parentHash = h
} else {
return nil, fmt.Errorf("expected uint64, got %T", hash)
}
}
event.ParentBlockHash = parentHash
if blocksize, err := parseInt64(data[4]); err == nil {
event.BlockSize = blocksize
} else {
return nil, fmt.Errorf("failed to parse field at index 4 as 'block_size': %w", err)
}
if medium, err := safeGetString(data[6]); err == nil {
event.Medium = medium
} else {
return nil, fmt.Errorf("failed to parse 'medium' from field at index 6: %w", err)
}
return event, nil
}
func parseVllmBlockRemoved(data []interface{}, timestamp interface{}) (*BlockRemovedEvent, error) {
event := &BlockRemovedEvent{
Type: EventTypeBlockRemoved,
}
for i, elem := range data {
slog.Debug("in parseVllmBlockRemoved:", "index", i, "type", fmt.Sprintf("%T", elem), "value", elem)
}
// Parse block hashes
if hashes, err := parseUint64Array(data[1]); err == nil {
event.BlockHashes = hashes
} else {
return nil, fmt.Errorf("failed to parse block_hashes: %w", err)
}
// parse medium
if medium, err := safeGetString(data[2]); err == nil {
event.Medium = medium
} else {
return nil, fmt.Errorf("failed to parse 'medium' from field at index 6: %w", err)
}
return event, nil
}
func convertToReplicaList(raw interface{}) ([][]string, error) {
list, ok := raw.([]interface{})
if !ok {
return nil, fmt.Errorf("expected []interface{}, got %T", raw)
}
result := make([][]string, len(list))
for i, item := range list {
subList, ok := item.([]interface{})
if !ok {
return nil, fmt.Errorf("item %d is not []interface{}, got %T", i, item)
}
result[i] = make([]string, len(subList))
for j, v := range subList {
str, ok := v.(string)
if !ok {
return nil, fmt.Errorf("element [%d][%d] is not string, got %T", i, j, v)
}
result[i][j] = str
}
}
return result, nil
}
func safeGetString(val interface{}) (string, error) {
switch v := val.(type) {
case string:
return v, nil
case []byte:
return string(v), nil
case fmt.Stringer:
return v.String(), nil
case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64, float32, float64:
return fmt.Sprintf("%v", v), nil
case nil:
return "", nil
default:
slog.Warn("Unexpected type in string field",
"type", fmt.Sprintf("%T", v),
"value", v)
return fmt.Sprintf("%v", v), nil
}
}
// Helper functions for parsing common types
func parseTimestamp(v interface{}) (time.Time, error) {
switch t := v.(type) {
case time.Time:
return t, nil
case int64:
return time.Unix(t, 0).UTC(), nil
case int:
return time.Unix(int64(t), 0).UTC(), nil
case int32:
return time.Unix(int64(t), 0).UTC(), nil
case uint32:
return time.Unix(int64(t), 0).UTC(), nil
case uint64:
return time.Unix(int64(t), 0).UTC(), nil
case float64:
sec := int64(t)
nsec := int64((t - float64(sec)) * 1e9)
return time.Unix(sec, nsec).UTC().Truncate(time.Microsecond), nil
case float32:
f64 := float64(t)
sec := int64(f64)
nsec := int64((f64 - float64(sec)) * 1e9)
return time.Unix(sec, nsec).UTC().Truncate(time.Microsecond), nil
case string:
// Try to parse RFC3339 format
return time.Parse(time.RFC3339, t)
default:
return time.Time{}, fmt.Errorf("unsupported timestamp type: %T", v)
}
}
func parseInt64(v interface{}) (int64, error) {
switch n := v.(type) {
case int64:
return n, nil
case int:
return int64(n), nil
case int32:
return int64(n), nil
case int16:
return int64(n), nil
case int8:
return int64(n), nil
case uint:
return int64(n), nil
case uint64:
return int64(n), nil
case uint32:
return int64(n), nil
case uint16:
return int64(n), nil
case uint8:
return int64(n), nil
case float64:
return int64(n), nil
case float32:
return int64(n), nil
default:
return 0, fmt.Errorf("unsupported int64 type: %T", v)
}
}
func parseUint64(v interface{}) (uint64, error) {
switch n := v.(type) {
case uint64:
return n, nil
case int:
return uint64(n), nil
case int32:
return uint64(n), nil
case int16:
return uint64(n), nil
case int8:
return uint64(n), nil
case uint:
return uint64(n), nil
case int64:
return uint64(n), nil
case uint32:
return uint64(n), nil
case uint16:
return uint64(n), nil
case uint8:
return uint64(n), nil
case float64:
return uint64(n), nil
case float32:
return uint64(n), nil
default:
return 0, fmt.Errorf("unsupported int64 type: %T", v)
}
}
func parseUint64Array(v interface{}) ([]uint64, error) {
arr, ok := v.([]interface{})
if !ok {
return nil, fmt.Errorf("expected array, got %T", v)
}
result := make([]uint64, 0, len(arr))
for i, item := range arr {
val, err := parseUint64(item)
if err != nil {
return nil, fmt.Errorf("failed to parse element at index %d: %w", i, err)
}
result = append(result, val)
}
return result, nil
}
func parseInt32Array(v interface{}) ([]int32, error) {
arr, ok := v.([]interface{})
if !ok {
return nil, fmt.Errorf("expected array, got %T", v)
}
result := make([]int32, 0, len(arr))
for i, item := range arr {
switch n := item.(type) {
case int32:
result = append(result, n)
case int:
result = append(result, int32(n))
case int64:
result = append(result, int32(n))
case int16:
result = append(result, int32(n))
case int8:
result = append(result, int32(n))
case uint:
result = append(result, int32(n))
case uint64:
result = append(result, int32(n))
case uint32:
result = append(result, int32(n))
case uint16:
result = append(result, int32(n))
case uint8:
result = append(result, int32(n))
case float64:
result = append(result, int32(n))
case float32:
result = append(result, int32(n))
default:
return nil, fmt.Errorf("unsupported int32 type at index %d: %T", i, item)
}
}
return result, nil
}
func parseMooncakeUint64(v interface{}) (uint64, error) {
var s string
switch val := v.(type) {
case string:
s = val
case nil:
s = ""
case uint64:
return val, nil
case int:
if val < 0 {
return 0, fmt.Errorf("negative value %d", val)
}
return uint64(val), nil
case int8:
if val < 0 {
return 0, fmt.Errorf("negative value %d", val)
}
return uint64(val), nil
case int16:
if val < 0 {
return 0, fmt.Errorf("negative value %d", val)
}
return uint64(val), nil
case int32:
if val < 0 {
return 0, fmt.Errorf("negative value %d", val)
}
return uint64(val), nil
case int64:
if val < 0 {
return 0, fmt.Errorf("negative value %d", val)
}
return uint64(val), nil
case uint:
return uint64(val), nil
case uint8:
return uint64(val), nil
case uint16:
return uint64(val), nil
case uint32:
return uint64(val), nil
default:
s = fmt.Sprint(v)
}
if s == "" {
return 0, nil
}
return strconv.ParseUint(s, 10, 64)
}
func parseMooncakeParentUint64(v interface{}) ([]uint64, error) {
switch val := v.(type) {
case nil:
return []uint64{}, nil
case string:
if val == "" {
return []uint64{}, nil
}
parts := strings.FieldsFunc(val, func(r rune) bool {
return r == ',' || r == ' ' || r == '\t' || r == '\n'
})
result := make([]uint64, 0, len(parts))
for _, part := range parts {
if part == "" {
continue
}
u, err := strconv.ParseUint(part, 10, 64)
if err != nil {
return nil, fmt.Errorf("failed to parse uint64 from string %q: %w", part, err)
}
result = append(result, u)
}
return result, nil
case []interface{}:
result := make([]uint64, 0, len(val))
for _, item := range val {
u, err := parseSingleUint64(item)
if err != nil {
return nil, fmt.Errorf("failed to parse element %v: %w", item, err)
}
result = append(result, u)
}
return result, nil
case []uint64:
// already correct type
return val, nil
default:
// try parse as single uint64
u, err := parseSingleUint64(val)
if err != nil {
return nil, err
}
return []uint64{u}, nil
}
}
func parseSingleUint64(v interface{}) (uint64, error) {
switch val := v.(type) {
case uint64:
return val, nil
case int:
if val < 0 {
return 0, fmt.Errorf("negative int %d cannot convert to uint64", val)
}
return uint64(val), nil
case int8:
if val < 0 {
return 0, fmt.Errorf("negative int8 %d cannot convert to uint64", val)
}
return uint64(val), nil
case int16:
if val < 0 {
return 0, fmt.Errorf("negative int16 %d cannot convert to uint64", val)
}
return uint64(val), nil
case int32:
if val < 0 {
return 0, fmt.Errorf("negative int32 %d cannot convert to uint64", val)
}
return uint64(val), nil
case int64:
if val < 0 {
return 0, fmt.Errorf("negative int64 %d cannot convert to uint64", val)
}
return uint64(val), nil
case uint:
return uint64(val), nil
case uint8:
return uint64(val), nil
case uint16:
return uint64(val), nil
case uint32:
return uint64(val), nil
case float32:
f := float64(val)
if f < 0 || f != float64(uint64(f)) {
return 0, fmt.Errorf("float32 %v invalid for uint64", val)
}
return uint64(f), nil
case float64:
if val < 0 || val != float64(uint64(val)) {
return 0, fmt.Errorf("float64 %v invalid for uint64", val)
}
return uint64(val), nil
case string:
if val == "" {
return 0, fmt.Errorf("empty string cannot be parsed as uint64")
}
return strconv.ParseUint(val, 10, 64)
case nil:
return 0, fmt.Errorf("nil cannot be parsed as uint64")
default:
return 0, fmt.Errorf("unsupported type %T for uint64 conversion", v)
}
}

View File

@ -0,0 +1,178 @@
package zmq_test
import (
"testing"
"time"
"conductor/zmq"
msgpack "github.com/shamaton/msgpack/v2"
)
func TestDecodeMooncakeEventBatch(t *testing.T) {
timestamp := int64(1700000000)
event := []interface{}{
"BlockStoreEvent",
"mooncake-key-123",
[][]interface{}{
[]interface{}{"replica1", "replica2"},
[]interface{}{"replica3"},
},
nil, // index 3 is not used
int64(1024), // BlockSize at index 4
[]interface{}{uint64(100), uint64(200)}, // BlockHashes at index 5
uint64(50), // ParentBlockHash at index 6
[]interface{}{int32(1), int32(2), int32(3)}, // TokenIDs at index 7
}
events := []interface{}{event}
batch := []interface{}{timestamp, events}
data, err := msgpack.Marshal(batch)
if err != nil {
t.Fatalf("Failed to marshal test data: %v", err)
}
result, err := zmq.DecodeMooncakeEventBatch(data)
if err != nil {
t.Fatalf("DecodeMooncakeEventBatch failed: %v", err)
}
if result.Source != zmq.SourceMooncake {
t.Errorf("Expected source %v, got %v", zmq.SourceMooncake, result.Source)
}
if len(result.Events) != 1 {
t.Fatalf("Expected 1 event, got %d", len(result.Events))
}
blockEvent, ok := result.Events[0].(*zmq.BlockStoredEvent)
if !ok {
t.Fatalf("Expected BlockStoredEvent, got %T", result.Events[0])
}
if blockEvent.Type != zmq.EventTypeBlockStored {
t.Errorf("Expected type %v, got %v", zmq.EventTypeBlockStored, blockEvent.Type)
}
if blockEvent.MooncakeKey != "mooncake-key-123" {
t.Errorf("Expected MooncakeKey 'mooncake-key-123', got '%s'", blockEvent.MooncakeKey)
}
if blockEvent.BlockSize != 1024 {
t.Errorf("Expected BlockSize 1024, got %d", blockEvent.BlockSize)
}
if len(blockEvent.BlockHashes) != 2 {
t.Fatalf("Expected 2 block hashes, got %d", len(blockEvent.BlockHashes))
}
if blockEvent.BlockHashes[0] != 100 || blockEvent.BlockHashes[1] != 200 {
t.Errorf("Expected block hashes [100, 200], got %v", blockEvent.BlockHashes)
}
if blockEvent.ParentBlockHash != 50 {
t.Errorf("Expected ParentBlockHash 50, got %d", blockEvent.ParentBlockHash)
}
if len(blockEvent.TokenIDs) != 3 {
t.Fatalf("Expected 3 token IDs, got %d", len(blockEvent.TokenIDs))
}
if len(blockEvent.ReplicaList) != 2 {
t.Fatalf("Expected 2 replica lists, got %d", len(blockEvent.ReplicaList))
}
}
func TestDecodeVllmEventBatch(t *testing.T) {
timestamp := int64(1700000000)
event := []interface{}{
"BlockStored",
[]interface{}{uint64(100), uint64(200)}, // BlockHashes
uint64(5000000000), // ParentBlockHash
[]interface{}{int32(10000000), int32(2), int32(3)}, // TokenIDs
int64(1024), // BlockSize
}
events := []interface{}{event}
status := "ok"
batch := []interface{}{timestamp, events, status}
data, err := msgpack.Marshal(batch)
if err != nil {
t.Fatalf("Failed to marshal test data: %v", err)
}
result, err := zmq.DecodeVllmEventBatch(data)
if err != nil {
t.Fatalf("DecodeVllmEventBatch failed: %v", err)
}
if result.Source != zmq.SourceVLLM {
t.Errorf("Expected source %v, got %v", zmq.SourceVLLM, result.Source)
}
if len(result.Events) != 1 {
t.Fatalf("Expected 1 event, got %d", len(result.Events))
}
blockEvent, ok := result.Events[0].(*zmq.BlockStoredEvent)
if !ok {
t.Fatalf("Expected BlockStoredEvent, got %T", result.Events[0])
}
if blockEvent.Type != zmq.EventTypeBlockStored {
t.Errorf("Expected type %v, got %v", zmq.EventTypeBlockStored, blockEvent.Type)
}
expectedTime := time.Unix(1700000000, 0).UTC()
if !blockEvent.Timestamp.Equal(expectedTime) {
t.Errorf("Expected timestamp %v, got %v", expectedTime, blockEvent.Timestamp)
}
if len(blockEvent.BlockHashes) != 2 {
t.Fatalf("Expected 2 block hashes, got %d", len(blockEvent.BlockHashes))
}
if blockEvent.BlockHashes[0] != 100 || blockEvent.BlockHashes[1] != 200 {
t.Errorf("Expected block hashes [100, 200], got %v", blockEvent.BlockHashes)
}
if blockEvent.ParentBlockHash != 5000000000 {
t.Errorf("Expected ParentBlockHash 50, got %d", blockEvent.ParentBlockHash)
}
if blockEvent.BlockSize != 1024 {
t.Errorf("Expected BlockSize 1024, got %d", blockEvent.BlockSize)
}
if len(blockEvent.TokenIDs) != 3 {
t.Fatalf("Expected 3 token IDs, got %d", len(blockEvent.TokenIDs))
}
}
func TestDecodeMooncakeEventBatch_InvalidData(t *testing.T) {
// Test with invalid array length
invalidBatch := []interface{}{int64(1700000000)} // Missing events
data, err := msgpack.Marshal(invalidBatch)
if err != nil {
t.Fatalf("Failed to marshal test data: %v", err)
}
_, err = zmq.DecodeMooncakeEventBatch(data)
if err == nil {
t.Error("Expected error for invalid array length, got nil")
}
}
func TestDecodeVllmEventBatch_InvalidData(t *testing.T) {
// Test with invalid array length
invalidBatch := []interface{}{int64(1700000000)} // Missing events and status
data, err := msgpack.Marshal(invalidBatch)
if err != nil {
t.Fatalf("Failed to marshal test data: %v", err)
}
_, err = zmq.DecodeVllmEventBatch(data)
if err == nil {
t.Error("Expected error for invalid array length, got nil")
}
}

View File

@ -0,0 +1,41 @@
package zmq
import (
"fmt"
"time"
)
// EventHandler processes received KV events
type EventHandler interface {
HandleEvent(event KVEvent, dpRank int64) error
}
// ZMQClientConfig contains configuration for the ZMQ client
type ZMQClientConfig struct {
CachePoolKey string
Endpoint string
ReplayEndpoint string
ModelName string
PollTimeout time.Duration
ReplayTimeout time.Duration
ReconnectDelay time.Duration
}
const (
// Timeouts and intervals
DefaultPollTimeout = 100 * time.Millisecond
DefaultReplayTimeout = 5 * time.Second
DefaultReconnectInterval = 1 * time.Second
MaxReconnectInterval = 30 * time.Second
ReconnectBackoffFactor = 2.0
EventChannelBufferSize = 1000
)
func ValidateConfig(config *ZMQClientConfig) error {
if config.Endpoint == "" {
return fmt.Errorf("endpoint is required")
}
return nil
}

View File

@ -0,0 +1,355 @@
package zmq
import (
"context"
"encoding/binary"
"fmt"
"log/slog"
"sync"
"time"
zmq "github.com/pebbe/zmq4"
)
type ZMQClient struct {
config *ZMQClientConfig
subSocket *zmq.Socket
replaySocket *zmq.Socket
eventHandler EventHandler
// State management
mu sync.RWMutex
connected bool
lastSeq int64
reconnectDelay time.Duration
// Lifecycle
ctx context.Context
cancel context.CancelFunc
wg sync.WaitGroup
}
func NewZMQClient(config *ZMQClientConfig, handler EventHandler) *ZMQClient {
ctx, cancel := context.WithCancel(context.Background())
return &ZMQClient{
config: config,
eventHandler: handler,
lastSeq: -1,
reconnectDelay: config.ReconnectDelay,
ctx: ctx,
cancel: cancel,
}
}
// Start initiates the connection and background event consumption loop.
func (c *ZMQClient) Start() error {
// Attempt initial connection
if err := c.Connect(); err != nil {
return fmt.Errorf("initial connection failed: %w", err)
}
c.wg.Add(1)
go c.loop()
slog.Info("ZMQ client started", "service", c.config.CachePoolKey)
return nil
}
func (c *ZMQClient) Stop() {
c.cancel()
c.wg.Wait()
c.mu.Lock()
c.cleanupSockets()
c.mu.Unlock()
slog.Info("ZMQ client stopped", "service", c.config.CachePoolKey)
}
// loop is the main background loop handling events and reconnections.
// Simplified: Fixed reconnect interval, single loop structure.
func (c *ZMQClient) loop() {
defer c.wg.Done()
for {
// Check if we should stop
select {
case <-c.ctx.Done():
return
default:
}
// 1. If disconnected, wait for ticker then try to reconnect
if !c.isConnected() {
c.handleReconnect()
continue
}
// 2. If connected, consume events
if err := c.consume(); err != nil {
slog.Error("Consumption error", "service", c.config.CachePoolKey, "error", err)
c.markDisconnected()
}
}
}
func (c *ZMQClient) handleReconnect() {
slog.Info("Attempting to reconnect to the service.", "service", c.config.CachePoolKey, "reconnectDelay", c.reconnectDelay)
ticker := time.NewTicker(c.config.ReconnectDelay)
defer ticker.Stop()
select {
case <-c.ctx.Done():
return
case <-ticker.C:
}
if err := c.Connect(); err != nil {
slog.Error("Reconnect failed", "service", c.config.CachePoolKey, "error", err)
}
// Reconnected! Request replay from last known sequence
lastSeq := c.getLastSequence()
if lastSeq >= 0 {
slog.Info("Reconnected", "service", c.config.CachePoolKey, "resuming_from", lastSeq+1)
if err := c.requestReplay(lastSeq + 1); err != nil {
slog.Warn("Failed to request replay after reconnect", "service", c.config.CachePoolKey, "error", err)
}
}
}
// Connect establishes the ZMQ SUB and DEALER sockets.
func (c *ZMQClient) Connect() error {
c.mu.Lock()
defer c.mu.Unlock()
if c.connected {
return nil
}
// Ensure clean state
c.cleanupSockets()
sock, err := zmq.NewSocket(zmq.SUB)
if err != nil {
return fmt.Errorf("create socket failed: %w", err)
}
if err := sock.SetIpv6(true); err != nil {
_ = sock.Close()
return fmt.Errorf("failed to enable IPv6 on socket: %w", err)
}
if err := sock.Connect(c.config.Endpoint); err != nil {
_ = sock.Close()
return fmt.Errorf("failed to connect to %s: %w", c.config.Endpoint, err)
}
// Important: Subscribe to all topics
if err := sock.SetSubscribe(""); err != nil {
_ = sock.Close()
return fmt.Errorf("failed to subscribe: %w", err)
}
replaySocket, err := zmq.NewSocket(zmq.DEALER)
if err != nil {
sock.Close()
return fmt.Errorf("failed to create DEALER socket: %w", err)
}
// Enable IPv6 for dual-stack support
if err := replaySocket.SetIpv6(true); err != nil {
_ = sock.Close()
_ = replaySocket.Close()
return fmt.Errorf("failed to enable IPv6 on DEALER socket: %w", err)
}
if err := replaySocket.Connect(c.config.ReplayEndpoint); err != nil {
_ = sock.Close()
_ = replaySocket.Close()
return fmt.Errorf("failed to connect to replay endpoint %s: %w", c.config.ReplayEndpoint, err)
}
c.subSocket = sock
c.replaySocket = replaySocket
c.connected = true
c.reconnectDelay = c.config.ReconnectDelay
slog.Info("Successfully connected to vLLM publisher", "service", c.config.CachePoolKey, "endpoint", c.config.Endpoint)
return nil
}
// consume reads and processes messages from the SUB socket.
func (c *ZMQClient) consume() error {
c.mu.RLock()
socket := c.subSocket
c.mu.RUnlock()
if socket == nil {
return fmt.Errorf("socket is nil")
}
poller := zmq.NewPoller()
poller.Add(socket, zmq.POLLIN)
// Poll for data
polled, err := poller.Poll(c.config.PollTimeout)
if err != nil {
return fmt.Errorf("poll error: %w", err)
}
if len(polled) == 0 {
return nil // No data, continue loop
}
if err := c.processMessage(socket); err != nil {
return fmt.Errorf("failed to process message: %w", err)
}
return nil
}
func (c *ZMQClient) processMessage(socket *zmq.Socket) error {
if socket == nil {
return fmt.Errorf("socket is nil")
}
// Read Frames: [Topic, Seq, Payload]
topic, err := socket.RecvBytes(0)
if err != nil {
return err
}
seqBytes, err := socket.RecvBytes(0)
if err != nil {
return err
}
payload, err := socket.RecvBytes(0)
if err != nil {
return err
}
if len(seqBytes) != 8 {
return fmt.Errorf("invalid sequence length")
}
seq := int64(binary.BigEndian.Uint64(seqBytes))
c.mu.RLock()
lastSeq := c.lastSeq
c.mu.RUnlock()
if lastSeq != -1 && seq > lastSeq+1 {
slog.Warn("Event gap detected",
"service", c.config.CachePoolKey,
"missed", seq-lastSeq-1,
"last", lastSeq,
"current", seq,
)
// Trigger replay for missed events?
// Usually we just log warning here, or could auto-trigger requestReplay
}
// Update Sequence immediately to keep state fresh
c.mu.Lock()
c.lastSeq = seq
c.mu.Unlock()
slog.Debug("enter deal topic", "topic", topic)
var batch *EventBatch
switch string(topic) {
case "mooncake":
batch, err = DecodeMooncakeEventBatch(payload)
default:
batch, err = DecodeVllmEventBatch(payload)
}
if err != nil {
return fmt.Errorf("decode failed: %w", err)
}
for _, event := range batch.Events {
// Inject Source Name
switch e := event.(type) {
case *BlockStoredEvent:
e.PodName = c.config.CachePoolKey
case *BlockRemovedEvent:
e.PodName = c.config.CachePoolKey
}
if err := c.eventHandler.HandleEvent(event, batch.DataParallelRank); err != nil {
slog.Error("Handler error", "service", c.config.CachePoolKey, "error", err)
}
}
slog.Debug("Processed batch", "service", c.config.CachePoolKey, "seq", seq, "topic", string(topic))
return nil
}
func (c *ZMQClient) requestReplay(fromSeq int64) error {
c.mu.RLock()
socket := c.replaySocket
c.mu.RUnlock()
if socket == nil {
return fmt.Errorf("replay socket is nil")
}
req := make([]byte, 8)
binary.BigEndian.PutUint64(req, uint64(fromSeq))
if _, err := socket.SendBytes(req, 0); err != nil {
return fmt.Errorf("failed to send replay request: %w", err)
}
// Ideally, we should wait for an ACK here if the protocol supports it
// For simplicity in static client, we fire and forget the request,
// assuming the server will send the replayed events via the SUB channel (or DEALER response)
// Original code read response from DEALER, let's keep that.
_ = socket.SetRcvtimeo(c.config.ReplayTimeout)
resp, err := socket.RecvBytes(0)
if err != nil {
return fmt.Errorf("failed to receive replay response: %w", err)
}
slog.Info("Replay requested", "service", c.config.CachePoolKey, "from", fromSeq, "resp_len", len(resp))
return nil
}
func (c *ZMQClient) cleanupSockets() {
if c.subSocket != nil {
c.subSocket.Close()
c.subSocket = nil
}
if c.replaySocket != nil {
c.replaySocket.Close()
c.replaySocket = nil
}
c.connected = false
}
func (c *ZMQClient) markDisconnected() {
c.mu.Lock()
defer c.mu.Unlock()
c.connected = false
}
func (c *ZMQClient) isConnected() bool {
c.mu.RLock()
defer c.mu.RUnlock()
return c.connected
}
func (c *ZMQClient) getLastSequence() int64 {
c.mu.RLock()
defer c.mu.RUnlock()
return c.lastSeq
}

View File

@ -0,0 +1,663 @@
package zmq_test
import (
"context"
"encoding/binary"
"errors"
"strconv"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
"conductor/zmq"
zmq4 "github.com/pebbe/zmq4"
msgpack "github.com/shamaton/msgpack/v2"
)
// MockEventHandler implements EventHandler for testing
type MockEventHandler struct {
mu sync.Mutex
events []zmq.KVEvent
handleError error
callCount int64
}
func NewMockEventHandler() *MockEventHandler {
return &MockEventHandler{
events: make([]zmq.KVEvent, 0),
}
}
func (m *MockEventHandler) HandleEvent(event zmq.KVEvent) error {
m.mu.Lock()
defer m.mu.Unlock()
atomic.AddInt64(&m.callCount, 1)
if m.handleError != nil {
return m.handleError
}
m.events = append(m.events, event)
return nil
}
func (m *MockEventHandler) GetEvents() []zmq.KVEvent {
m.mu.Lock()
defer m.mu.Unlock()
events := make([]zmq.KVEvent, len(m.events))
copy(events, m.events)
return events
}
func (m *MockEventHandler) GetCallCount() int64 {
return atomic.LoadInt64(&m.callCount)
}
func (m *MockEventHandler) SetHandleError(err error) {
m.mu.Lock()
defer m.mu.Unlock()
m.handleError = err
}
func (m *MockEventHandler) Clear() {
m.mu.Lock()
defer m.mu.Unlock()
m.events = m.events[:0]
m.handleError = nil
atomic.StoreInt64(&m.callCount, 0)
}
// MockPublisher simulates a ZMQ publisher for testing
type MockPublisher struct {
pubSocket *zmq4.Socket
routerSocket *zmq4.Socket
sequence int64
mu sync.Mutex
ctx context.Context
cancel context.CancelFunc
wg sync.WaitGroup
}
func createMockPublisher(t *testing.T, pubPort, routerPort int) *MockPublisher {
ctx, cancel := context.WithCancel(context.Background())
// Create PUB socket
pubSocket, err := zmq4.NewSocket(zmq4.PUB)
if err != nil {
t.Fatalf("Failed to create PUB socket: %v", err)
}
err = pubSocket.SetIpv6(true)
if err != nil {
pubSocket.Close()
t.Fatalf("Failed to enable IPv6 on PUB socket: %v", err)
}
err = pubSocket.Bind("tcp://127.0.0.1:*")
if err != nil {
pubSocket.Close()
t.Fatalf("Failed to bind PUB socket: %v", err)
}
// Create ROUTER socket for replay
routerSocket, err := zmq4.NewSocket(zmq4.ROUTER)
if err != nil {
pubSocket.Close()
t.Fatalf("Failed to create ROUTER socket: %v", err)
}
err = routerSocket.SetIpv6(true)
if err != nil {
pubSocket.Close()
routerSocket.Close()
t.Fatalf("Failed to enable IPv6 on ROUTER socket: %v", err)
}
err = routerSocket.Bind("tcp://127.0.0.1:*")
if err != nil {
pubSocket.Close()
routerSocket.Close()
t.Fatalf("Failed to bind ROUTER socket: %v", err)
}
mp := &MockPublisher{
pubSocket: pubSocket,
routerSocket: routerSocket,
ctx: ctx,
cancel: cancel,
}
// Start replay handler
mp.wg.Add(1)
go mp.handleReplay()
// Wait for sockets to bind
time.Sleep(100 * time.Millisecond)
return mp
}
func (mp *MockPublisher) PublishEvent(topic string, event zmq.KVEvent) error {
// Encode event based on topic
var payload []byte
var err error
if topic == "mooncake" {
payload, err = encodeMooncakeEvent(event)
} else {
payload, err = encodeVllmEvent(event)
}
if err != nil {
return err
}
mp.mu.Lock()
mp.sequence++
seq := mp.sequence
mp.mu.Unlock()
seqBytes := make([]byte, 8)
binary.BigEndian.PutUint64(seqBytes, uint64(seq))
_, err = mp.pubSocket.SendMessage(topic, seqBytes, payload)
return err
}
func (mp *MockPublisher) handleReplay() {
defer mp.wg.Done()
for {
select {
case <-mp.ctx.Done():
return
default:
}
// Set receive timeout
_ = mp.routerSocket.SetRcvtimeo(100 * time.Millisecond)
msg, err := mp.routerSocket.RecvBytes(0)
if err != nil {
continue
}
if len(msg) == 8 {
// Send ACK
_, _ = mp.routerSocket.SendBytes([]byte("OK"), 0)
}
}
}
func (mp *MockPublisher) Close() {
mp.cancel()
mp.wg.Wait()
_ = mp.pubSocket.Close()
_ = mp.routerSocket.Close()
}
func encodeMooncakeEvent(event zmq.KVEvent) ([]byte, error) {
switch e := event.(type) {
case *zmq.BlockStoredEvent:
timestamp := e.Timestamp.Unix()
eventData := []interface{}{
"BlockStoreEvent",
e.MooncakeKey,
e.ReplicaList,
nil, // index 3 not used
e.BlockSize,
convertUint64Slice(e.BlockHashes),
e.ParentBlockHash,
convertInt32Slice(e.TokenIDs),
}
batch := []interface{}{timestamp, []interface{}{eventData}}
return msgpack.Marshal(batch)
default:
return nil, errors.New("unsupported event type for mooncake")
}
}
func encodeVllmEvent(event zmq.KVEvent) ([]byte, error) {
switch e := event.(type) {
case *zmq.BlockStoredEvent:
timestamp := e.Timestamp.Unix()
eventData := []interface{}{
"BlockStored",
convertUint64Slice(e.BlockHashes),
e.ParentBlockHash,
convertInt32Slice(e.TokenIDs),
e.BlockSize,
}
batch := []interface{}{timestamp, []interface{}{eventData}, "ok"}
return msgpack.Marshal(batch)
case *zmq.BlockRemovedEvent:
timestamp := e.Timestamp.Unix()
eventData := []interface{}{
"BlockRemoved",
convertUint64Slice(e.BlockHashes),
}
batch := []interface{}{timestamp, []interface{}{eventData}, "ok"}
return msgpack.Marshal(batch)
default:
return nil, errors.New("unsupported event type for vllm")
}
}
func convertUint64Slice(slice []uint64) []interface{} {
result := make([]interface{}, len(slice))
for i, v := range slice {
result[i] = uint64(v)
}
return result
}
func convertInt32Slice(slice []int32) []interface{} {
result := make([]interface{}, len(slice))
for i, v := range slice {
result[i] = int32(v)
}
return result
}
func skipIfZMQUnavailable(t *testing.T) {
ctx, err := zmq4.NewContext()
if err != nil {
t.Skip("ZMQ not available:", err)
}
_ = ctx.Term()
}
func TestZMQClient_Connect_Success(t *testing.T) {
skipIfZMQUnavailable(t)
publisher := createMockPublisher(t, 5547, 5548)
defer publisher.Close()
// Get actual bound ports
pubEndpoint, err := publisher.pubSocket.GetLastEndpoint()
if err != nil {
t.Fatalf("Failed to get PUB endpoint: %v", err)
}
routerEndpoint, err := publisher.routerSocket.GetLastEndpoint()
if err != nil {
t.Fatalf("Failed to get ROUTER endpoint: %v", err)
}
pubPort := extractPortFromEndpoint(pubEndpoint)
routerPort := extractPortFromEndpoint(routerEndpoint)
config := &zmq.ZMQClientConfig{
CachePoolKey: "test-pod",
ServiceIP: "127.0.0.1",
ModelName: "test-model",
Port: pubPort,
RouterPort: routerPort,
PollTimeout: 100 * time.Millisecond,
ReplayTimeout: 1 * time.Second,
ReconnectDelay: 100 * time.Millisecond,
}
handler := NewMockEventHandler()
client := zmq.NewZMQClient(config, handler)
err = client.Connect()
if err != nil {
t.Fatalf("Connect failed: %v", err)
}
client.Stop()
}
func TestZMQClient_Connect_AlreadyConnected(t *testing.T) {
skipIfZMQUnavailable(t)
publisher := createMockPublisher(t, 5557, 5558)
defer publisher.Close()
pubEndpoint, _ := publisher.pubSocket.GetLastEndpoint()
routerEndpoint, _ := publisher.routerSocket.GetLastEndpoint()
pubPort := extractPortFromEndpoint(pubEndpoint)
routerPort := extractPortFromEndpoint(routerEndpoint)
config := &zmq.ZMQClientConfig{
CachePoolKey: "test-pod",
ServiceIP: "127.0.0.1",
ModelName: "test-model",
Port: pubPort,
RouterPort: routerPort,
PollTimeout: 100 * time.Millisecond,
ReplayTimeout: 1 * time.Second,
ReconnectDelay: 100 * time.Millisecond,
}
handler := NewMockEventHandler()
client := zmq.NewZMQClient(config, handler)
err := client.Connect()
if err != nil {
t.Fatalf("First Connect failed: %v", err)
}
// Connect again should not error
err = client.Connect()
if err != nil {
t.Fatalf("Second Connect failed: %v", err)
}
client.Stop()
}
func TestZMQClient_Start_Stop(t *testing.T) {
skipIfZMQUnavailable(t)
publisher := createMockPublisher(t, 5557, 5558)
defer publisher.Close()
pubEndpoint, _ := publisher.pubSocket.GetLastEndpoint()
routerEndpoint, _ := publisher.routerSocket.GetLastEndpoint()
pubPort := extractPortFromEndpoint(pubEndpoint)
routerPort := extractPortFromEndpoint(routerEndpoint)
config := &zmq.ZMQClientConfig{
CachePoolKey: "test-pod",
ServiceIP: "127.0.0.1",
ModelName: "test-model",
Port: pubPort,
RouterPort: routerPort,
PollTimeout: 100 * time.Millisecond,
ReplayTimeout: 1 * time.Second,
ReconnectDelay: 100 * time.Millisecond,
}
handler := NewMockEventHandler()
client := zmq.NewZMQClient(config, handler)
err := client.Start()
if err != nil {
t.Fatalf("Start failed: %v", err)
}
// Wait a bit for loop to start
time.Sleep(50 * time.Millisecond)
// Stop should work gracefully
client.Stop()
}
func TestZMQClient_ProcessMessage_MooncakeTopic(t *testing.T) {
skipIfZMQUnavailable(t)
publisher := createMockPublisher(t, 5557, 5558)
defer publisher.Close()
pubEndpoint, _ := publisher.pubSocket.GetLastEndpoint()
routerEndpoint, _ := publisher.routerSocket.GetLastEndpoint()
pubPort := extractPortFromEndpoint(pubEndpoint)
routerPort := extractPortFromEndpoint(routerEndpoint)
config := &zmq.ZMQClientConfig{
CachePoolKey: "test-pod",
ServiceIP: "127.0.0.1",
ModelName: "test-model",
Port: pubPort,
RouterPort: routerPort,
PollTimeout: 100 * time.Millisecond,
ReplayTimeout: 1 * time.Second,
ReconnectDelay: 100 * time.Millisecond,
}
handler := NewMockEventHandler()
client := zmq.NewZMQClient(config, handler)
err := client.Start()
if err != nil {
t.Fatalf("Start failed: %v", err)
}
defer client.Stop()
// Wait for connection
time.Sleep(100 * time.Millisecond)
// Create and publish mooncake event
event := &zmq.BlockStoredEvent{
Type: zmq.EventTypeBlockStored,
Timestamp: time.Now().UTC(),
BlockHashes: []uint64{100, 200},
TokenIDs: []int32{1, 2, 3},
ParentBlockHash: 50,
BlockSize: 1024,
MooncakeKey: "mooncake-key-123",
ReplicaList: [][]string{{"replica1", "replica2"}},
ModelName: "test-model",
}
err = publisher.PublishEvent("mooncake", event)
if err != nil {
t.Fatalf("PublishEvent failed: %v", err)
}
// Wait for processing
time.Sleep(200 * time.Millisecond)
events := handler.GetEvents()
if len(events) < 1 {
t.Fatalf("Expected at least 1 event, got %d", len(events))
}
blockEvent, ok := events[0].(*zmq.BlockStoredEvent)
if !ok {
t.Fatalf("Expected BlockStoredEvent, got %T", events[0])
}
if blockEvent.PodName != "test-pod" {
t.Errorf("Expected PodName 'test-pod', got '%s'", blockEvent.PodName)
}
if blockEvent.MooncakeKey != "mooncake-key-123" {
t.Errorf("Expected MooncakeKey 'mooncake-key-123', got '%s'", blockEvent.MooncakeKey)
}
if len(blockEvent.BlockHashes) == 0 || blockEvent.BlockHashes[0] != 100 {
t.Errorf("Expected first BlockHash 100, got %v", blockEvent.BlockHashes)
}
}
func TestZMQClient_ProcessMessage_VllmTopic(t *testing.T) {
skipIfZMQUnavailable(t)
publisher := createMockPublisher(t, 5557, 5558)
defer publisher.Close()
pubEndpoint, _ := publisher.pubSocket.GetLastEndpoint()
routerEndpoint, _ := publisher.routerSocket.GetLastEndpoint()
pubPort := extractPortFromEndpoint(pubEndpoint)
routerPort := extractPortFromEndpoint(routerEndpoint)
config := &zmq.ZMQClientConfig{
CachePoolKey: "test-pod",
ServiceIP: "127.0.0.1",
ModelName: "test-model",
Port: pubPort,
RouterPort: routerPort,
PollTimeout: 100 * time.Millisecond,
ReplayTimeout: 1 * time.Second,
ReconnectDelay: 100 * time.Millisecond,
}
handler := NewMockEventHandler()
client := zmq.NewZMQClient(config, handler)
err := client.Start()
if err != nil {
t.Fatalf("Start failed: %v", err)
}
defer client.Stop()
// Wait for connection
time.Sleep(100 * time.Millisecond)
// Create and publish vllm event
event := &zmq.BlockStoredEvent{
Type: zmq.EventTypeBlockStored,
Timestamp: time.Now().UTC(),
BlockHashes: []uint64{300, 400},
TokenIDs: []int32{4, 5, 6},
ParentBlockHash: 1500000000000000,
BlockSize: 2048,
ModelName: "test-model",
}
err = publisher.PublishEvent("vllm", event)
if err != nil {
t.Fatalf("PublishEvent failed: %v", err)
}
// Wait for processing
time.Sleep(200 * time.Millisecond)
events := handler.GetEvents()
if len(events) < 1 {
t.Fatalf("Expected at least 1 event, got %d", len(events))
}
blockEvent, ok := events[0].(*zmq.BlockStoredEvent)
if !ok {
t.Fatalf("Expected BlockStoredEvent, got %T", events[0])
}
if blockEvent.PodName != "test-pod" {
t.Errorf("Expected PodName 'test-pod', got '%s'", blockEvent.PodName)
}
if len(blockEvent.BlockHashes) == 0 || blockEvent.BlockHashes[0] != 300 {
t.Errorf("Expected first BlockHash 300, got %v", blockEvent.BlockHashes)
}
}
func TestZMQClient_SequenceTracking(t *testing.T) {
skipIfZMQUnavailable(t)
publisher := createMockPublisher(t, 5557, 5558)
defer publisher.Close()
pubEndpoint, _ := publisher.pubSocket.GetLastEndpoint()
routerEndpoint, _ := publisher.routerSocket.GetLastEndpoint()
pubPort := extractPortFromEndpoint(pubEndpoint)
routerPort := extractPortFromEndpoint(routerEndpoint)
config := &zmq.ZMQClientConfig{
CachePoolKey: "test-pod",
ServiceIP: "127.0.0.1",
ModelName: "test-model",
Port: pubPort,
RouterPort: routerPort,
PollTimeout: 100 * time.Millisecond,
ReplayTimeout: 1 * time.Second,
ReconnectDelay: 100 * time.Millisecond,
}
handler := NewMockEventHandler()
client := zmq.NewZMQClient(config, handler)
err := client.Start()
if err != nil {
t.Fatalf("Start failed: %v", err)
}
defer client.Stop()
// Wait for connection
time.Sleep(100 * time.Millisecond)
// Publish multiple events
for i := 0; i < 5; i++ {
event := &zmq.BlockStoredEvent{
Type: zmq.EventTypeBlockStored,
Timestamp: time.Now().UTC(),
BlockHashes: []uint64{uint64(i)},
TokenIDs: []int32{int32(i)},
ParentBlockHash: 1000000000000000000,
BlockSize: 128,
ModelName: "test-model",
}
_ = publisher.PublishEvent("vllm", event)
time.Sleep(50 * time.Millisecond)
}
// Wait for processing
time.Sleep(500 * time.Millisecond)
// Verify events were processed
events := handler.GetEvents()
if len(events) < 5 {
t.Errorf("Expected at least 5 events, got %d", len(events))
}
}
func TestZMQClient_ProcessMessage_EventGap(t *testing.T) {
skipIfZMQUnavailable(t)
publisher := createMockPublisher(t, 5557, 5558)
defer publisher.Close()
pubEndpoint, _ := publisher.pubSocket.GetLastEndpoint()
routerEndpoint, _ := publisher.routerSocket.GetLastEndpoint()
pubPort := extractPortFromEndpoint(pubEndpoint)
routerPort := extractPortFromEndpoint(routerEndpoint)
config := &zmq.ZMQClientConfig{
CachePoolKey: "test-pod",
ServiceIP: "127.0.0.1",
ModelName: "test-model",
Port: pubPort,
RouterPort: routerPort,
PollTimeout: 100 * time.Millisecond,
ReplayTimeout: 1 * time.Second,
ReconnectDelay: 100 * time.Millisecond,
}
handler := NewMockEventHandler()
client := zmq.NewZMQClient(config, handler)
err := client.Start()
if err != nil {
t.Fatalf("Start failed: %v", err)
}
defer client.Stop()
// Wait for connection
time.Sleep(100 * time.Millisecond)
// Manually set last sequence to simulate gap
// This is a bit tricky since we need to access internal state
// For now, we publish events and verify gap detection works
// The actual gap detection is logged, so we verify the code path exists
event := &zmq.BlockStoredEvent{
Type: zmq.EventTypeBlockStored,
Timestamp: time.Now().UTC(),
BlockHashes: []uint64{100},
TokenIDs: []int32{1},
ModelName: "test-model",
}
_ = publisher.PublishEvent("vllm", event)
time.Sleep(100 * time.Millisecond)
// Publish another event - gap detection should work for subsequent events
_ = publisher.PublishEvent("vllm", event)
time.Sleep(100 * time.Millisecond)
}
// Helper function to extract port from endpoint string like "tcp://127.0.0.1:5557"
func extractPortFromEndpoint(endpoint string) int {
// Parse "tcp://127.0.0.1:5557" to get 5557
parts := strings.Split(endpoint, ":")
if len(parts) < 3 {
return 5557 // Default fallback
}
portStr := parts[len(parts)-1]
port, err := strconv.Atoi(portStr)
if err != nil {
return 5557 // Default fallback
}
return port
}

View File

@ -0,0 +1,355 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Adapted from the vLLM repository's toy_proxy_server.py in tests/v1/kv_connector/nixl_integration/.
import argparse
import itertools
import logging
import os
import uuid
from contextlib import asynccontextmanager
import httpx
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
class CacheAwareRouter():
def __init__(self, address, endpoint):
self.address = address
self.endpoint = endpoint
self.client = httpx.AsyncClient(timeout=None, base_url=f'http://{address}')
async def get_best_prefiller(self, token_ids: list, ready_instances, req_data):
# call conductor restful api to get cache hit situation
model_name = req_data.get("model", "ds")
lora_id = req_data.get("lora_id", -1)
request_data = {
"instances": ready_instances,
"token_ids": token_ids,
"model_name": model_name,
"lora_id": lora_id
}
headers = {
"Content-Type": "application/json",
}
logger.debug(f"conductor request_data: {request_data}")
response = await self.client.post(self.endpoint, json=request_data, headers=headers)
response.raise_for_status()
return response.json()["HitStatus"]
async def close(self) -> None:
if self.client:
await self.client.aclose()
@asynccontextmanager
async def lifespan(app: FastAPI):
"""
Lifespan context manager to handle startup and shutdown events.
"""
# Startup: Initialize client pools for prefiller and decoder services
app.state.prefill_clients = []
app.state.decode_clients = []
# Create prefill clients
for i, (host, port) in enumerate(global_args.prefiller_instances):
prefiller_base_url = f'http://{host}:{port}'
app.state.prefill_clients.append({
'client':
httpx.AsyncClient(timeout=None, base_url=prefiller_base_url),
'host':
host,
'port':
port,
'id':
i
})
# Create decode clients
for i, (host, port) in enumerate(global_args.decoder_instances):
decoder_base_url = f'http://{host}:{port}'
app.state.decode_clients.append({
'client':
httpx.AsyncClient(timeout=None, base_url=decoder_base_url),
'host':
host,
'port':
port,
'id':
i
})
# Create conductor client
app.state.conductor_client = CacheAwareRouter(global_args.conductor_address, "/cache")
# Initialize round-robin iterators
app.state.prefill_iterator = itertools.cycle(
range(len(app.state.prefill_clients)))
app.state.decode_iterator = itertools.cycle(
range(len(app.state.decode_clients)))
logger.info(f"Initialized {len(app.state.prefill_clients)} prefill clients "
f"and {len(app.state.decode_clients)} decode clients.")
yield
# Shutdown: Close all clients
for client_info in app.state.prefill_clients:
await client_info['client'].aclose()
for client_info in app.state.decode_clients:
await client_info['client'].aclose()
# Close conductor client
await app.state.conductor_client.close()
# Update FastAPI app initialization to use lifespan
app = FastAPI(lifespan=lifespan)
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--host", type=str, default="localhost")
# For prefiller instances
parser.add_argument("--prefiller-hosts",
"--prefiller-host",
type=str,
nargs="+",
default=["localhost"])
parser.add_argument("--prefiller-ports",
"--prefiller-port",
type=int,
nargs="+",
default=[8100])
# For decoder instances
parser.add_argument("--decoder-hosts",
"--decoder-host",
type=str,
nargs="+",
default=["localhost"])
parser.add_argument("--decoder-ports",
"--decoder-port",
type=int,
nargs="+",
default=[8200])
parser.add_argument("--conductor-address", type=str, default="127.0.0.1:13333")
args = parser.parse_args()
# Validate and pair hosts with ports
if len(args.prefiller_hosts) != len(args.prefiller_ports):
raise ValueError(
"Number of prefiller hosts must match number of prefiller ports")
if len(args.decoder_hosts) != len(args.decoder_ports):
raise ValueError(
"Number of decoder hosts must match number of decoder ports")
# Create tuples of (host, port) for each service type
args.prefiller_instances = list(
zip(args.prefiller_hosts, args.prefiller_ports))
args.decoder_instances = list(zip(args.decoder_hosts, args.decoder_ports))
return args
def get_next_client(app, service_type: str):
"""
Get the next client in round-robin fashion.
Args:
app: The FastAPI app instance
service_type: Either 'prefill' or 'decode'
Returns:
The next client to use
"""
if service_type == 'prefill':
client_idx = next(app.state.prefill_iterator)
return app.state.prefill_clients[client_idx]
elif service_type == 'decode':
client_idx = next(app.state.decode_iterator)
return app.state.decode_clients[client_idx]
else:
raise ValueError(f"Unknown service type: {service_type}")
async def get_best_prefiller(app, token_ids: list, round_robin_prefill, req_data):
# Get all prefill instances
ready_instances = []
index_map = {}
for index, client_info in enumerate(app.state.prefill_clients):
if client_info['client'].is_closed:
continue
ready_instances.append(client_info['host'])
index_map[client_info['host']] = index
cache_hit_status = await app.state.conductor_client.get_best_prefiller(token_ids, ready_instances, req_data)
if not cache_hit_status:
return round_robin_prefill
best_prefiller_index = None
max_hit_value = -1
for k, v in cache_hit_status.items():
if v > max_hit_value:
best_prefiller_index = index_map[k]
max_hit_value = v
return app.state.prefill_clients[best_prefiller_index]
async def get_tokenid(client_info: dict, req_data: dict, request_id: str):
req_data = req_data.copy()
req_data["stream"] = False
req_data["max_tokens"] = 1
if "stream_options" in req_data:
del req_data["stream_options"]
headers = {
"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
"X-Request-Id": request_id
}
response = await client_info['client'].post("/tokenize",
json=req_data,
headers=headers)
response.raise_for_status()
token_id = response.json()["tokens"]
return token_id
async def send_request_to_service(client_info: dict, endpoint: str,
req_data: dict, request_id: str):
"""
Send a request to a service using a client from the pool.
"""
req_data = req_data.copy()
req_data['kv_transfer_params'] = {
"do_remote_decode": True,
"do_remote_prefill": False,
"remote_engine_id": None,
"remote_block_ids": None,
"remote_host": None,
"remote_port": None
}
req_data["stream"] = False
req_data["max_tokens"] = 1
if "max_completion_tokens" in req_data:
req_data["max_completion_tokens"] = 1
if "stream_options" in req_data:
del req_data["stream_options"]
headers = {
"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
"X-Request-Id": request_id
}
logger.debug(f"req_data: {req_data}")
response = await client_info['client'].post(endpoint,
json=req_data,
headers=headers)
response.raise_for_status()
return response
async def stream_service_response(client_info: dict, endpoint: str,
req_data: dict, request_id: str):
"""
Asynchronously stream response from a service using a client from the pool.
"""
headers = {
"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
"X-Request-Id": request_id
}
async with client_info['client'].stream("POST",
endpoint,
json=req_data,
headers=headers) as response:
response.raise_for_status()
async for chunk in response.aiter_bytes():
yield chunk
async def _handle_completions(api: str, request: Request):
try:
req_data = await request.json()
request_id = str(uuid.uuid4())
# select tokenizer client in round-robin fashion
remote_tokenizer_client_info = get_next_client(request.app, 'prefill')
token_ids = await get_tokenid(remote_tokenizer_client_info, req_data, request_id)
# choice best cache hit prefill instance
prefill_client_info = await get_best_prefiller(request.app, token_ids, remote_tokenizer_client_info, req_data)
response = await send_request_to_service(prefill_client_info, api,
req_data, request_id)
# Extract the needed fields
response_json = response.json()
kv_transfer_params = response_json.get('kv_transfer_params', {})
if kv_transfer_params:
req_data["kv_transfer_params"] = kv_transfer_params
# Get the next decode client in round-robin fashion
decode_client_info = get_next_client(request.app, 'decode')
logger.debug("Using %s %s", prefill_client_info, decode_client_info)
# Stream response from decode service
async def generate_stream():
async for chunk in stream_service_response(decode_client_info,
api,
req_data,
request_id=request_id):
yield chunk
return StreamingResponse(generate_stream(),
media_type="application/json")
except Exception as e:
import sys
import traceback
exc_info = sys.exc_info()
logger.error("Error occurred in disagg prefill proxy server"
f" - {api} endpoint")
logger.error(e)
logger.error("".join(traceback.format_exception(*exc_info)))
raise
@app.post("/v1/completions")
async def handle_completions(request: Request):
return await _handle_completions("/v1/completions", request)
@app.post("/v1/chat/completions")
async def handle_chat_completions(request: Request):
return await _handle_completions("/v1/chat/completions", request)
@app.get("/healthcheck")
async def healthcheck():
"""Simple endpoint to check if the server is running."""
return {
"status": "ok",
"prefill_instances": len(app.state.prefill_clients),
"decode_instances": len(app.state.decode_clients)
}
if __name__ == '__main__':
global global_args
global_args = parse_args()
import uvicorn
uvicorn.run(app, host=global_args.host, port=global_args.port)

View File

@ -0,0 +1,32 @@
{
"kvevent_instance":
{
"vllm-prefill-node1":
{
"endpoint": "tcp://127.0.0.1:5557",
"replay_endpoint": "tcp://127.0.0.1:5558",
"type": "vLLM",
"modelname": "qwen2.5",
"lora_name": "xx-adapter",
"tenant_id": "default",
"instance_id": "vllm-prefill-node1",
"block_size": 128,
"dp_rank": 0,
"additionalsalt": ""
},
"mooncake":
{
"endpoint": "tcp://127.0.0.1:6667",
"replay_endpoint": "tcp://127.0.0.1:6668",
"type": "Mooncake",
"modelname": "qwen2.5",
"lora_name": "xx-adapter",
"tenant_id": "default",
"instance_id": "vllm-prefill-node1",
"block_size": 128,
"dp_rank": 0,
"additionalsalt": ""
}
},
"http_server_port": 13333
}

View File

@ -1,93 +0,0 @@
# BuildEpExt.cmake - Build the Mooncake EP Python extension.
#
# Invoked at build time via cmake -P from the root CMakeLists.txt when
# WITH_EP=ON. Variables are passed with -D from the custom target:
#
# SOURCE_DIR - mooncake-ep source directory
# EP_CUDA_MAJOR - CUDA major version (integer)
# EP_TORCH_VERSIONS - pipe-separated (|) PyTorch versions to build for
# (empty = use the currently-installed torch)
# TORCH_CUDA_ARCH_LIST - pipe-separated CUDA arch list forwarded to torch
# STAGING_DIR - destination directory for the built .so files
# ENGINE_SO_PATH - absolute path to the built engine.cpython-XYZ.so
cmake_minimum_required(VERSION 3.16)
# Include common build utilities.
include("${SOURCE_DIR}/../mooncake-common/SetupPyTorchEnv.cmake")
# Restore pipe-separated strings back to CMake semicolon-separated lists.
if(EP_TORCH_VERSIONS)
string(REPLACE "|" ";" EP_TORCH_VERSIONS "${EP_TORCH_VERSIONS}")
endif()
if(TORCH_CUDA_ARCH_LIST)
string(REPLACE "|" ";" TORCH_CUDA_ARCH_LIST "${TORCH_CUDA_ARCH_LIST}")
endif()
# ---------------------------------------------------------------------------
# 1. Set up the build environment.
# ---------------------------------------------------------------------------
# Clear jobserver variables so that sub-processes started by setup.py do not
# try to connect to the parent ninja's jobserver pipe FDs, which are not
# inherited and cause: "ninja: error: Could not initialize jobserver: Invalid
# file descriptors".
set(ENV{MAKEFLAGS} "")
set(ENV{MFLAGS} "")
set(ENV{TORCH_CUDA_ARCH_LIST} "${TORCH_CUDA_ARCH_LIST}")
# ---------------------------------------------------------------------------
# 2. Ensure engine.so exists in mooncake-wheel/mooncake/ for setup.py linking.
# ---------------------------------------------------------------------------
# setup.py links against -l:engine.so in ../mooncake-wheel/mooncake/.
# During the make phase only the versioned engine.cpython-XYZ.so exists in
# the build tree; create a bare engine.so symlink so the linker can find it.
set(_wheel_mooncake_dir "${SOURCE_DIR}/../mooncake-wheel/mooncake")
set(_engine_symlink "${_wheel_mooncake_dir}/engine.so")
if(ENGINE_SO_PATH AND NOT EXISTS "${_engine_symlink}")
message(STATUS "[EP] Creating engine.so symlink -> ${ENGINE_SO_PATH}")
execute_process(
COMMAND ${CMAKE_COMMAND} -E create_symlink "${ENGINE_SO_PATH}" "${_engine_symlink}"
)
endif()
# ---------------------------------------------------------------------------
# 3. Build the EP Python extension.
# ---------------------------------------------------------------------------
if("${EP_TORCH_VERSIONS}" STREQUAL "")
message(STATUS "[EP] Building with currently-installed PyTorch")
execute_process(
COMMAND ${Python3_EXECUTABLE} setup.py build_ext --build-lib .
WORKING_DIRECTORY "${SOURCE_DIR}"
RESULT_VARIABLE _ret
)
if(NOT _ret EQUAL 0)
message(FATAL_ERROR "[EP] Extension build failed (exit code: ${_ret})")
endif()
else()
message(STATUS "[EP] Building for PyTorch versions: ${EP_TORCH_VERSIONS}")
foreach(_version IN LISTS EP_TORCH_VERSIONS)
install_pytorch_wheel("${_version}" "${EP_CUDA_MAJOR}" "${EP_CUDA_MINOR}" "[EP]")
execute_process(
COMMAND ${Python3_EXECUTABLE} setup.py build_ext --build-lib . --force
WORKING_DIRECTORY "${SOURCE_DIR}"
RESULT_VARIABLE _ret
)
if(NOT _ret EQUAL 0)
message(FATAL_ERROR "[EP] Extension build failed for PyTorch ${_version}")
endif()
endforeach()
endif()
# ---------------------------------------------------------------------------
# 4. Copy the built .so files to the staging directory.
# ---------------------------------------------------------------------------
file(MAKE_DIRECTORY "${STAGING_DIR}")
file(GLOB _so_files "${SOURCE_DIR}/mooncake/*.so")
foreach(_so IN LISTS _so_files)
get_filename_component(_fname "${_so}" NAME)
message(STATUS "[EP] Staging ${_fname} -> ${STAGING_DIR}")
file(COPY "${_so}" DESTINATION "${STAGING_DIR}" NO_SOURCE_PERMISSIONS)
endforeach()
message(STATUS "[EP] Mooncake EP extension build complete")

View File

@ -30,4 +30,13 @@ find_package(Torch REQUIRED)
include_directories(${TORCH_INCLUDE_DIRS})
include_directories(include)
add_subdirectory(include)
add_subdirectory(src)
if (BUILD_UNIT_TESTS)
add_subdirectory(tests)
endif()
if (BUILD_EXAMPLES)
add_subdirectory(example)
endif()

View File

View File

@ -3,7 +3,6 @@
#include <ATen/cuda/CUDAContext.h>
#include <cuda_bf16.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <fstream>
#include <mooncake_ibgda/memheap.h>
@ -72,7 +71,7 @@ struct MooncakeEpBuffer {
void* gdr_buffer = nullptr;
// IBGDA
static constexpr size_t CTRL_BUF_SIZE = 1024ULL * 1024 * 1024; // 1024 MiB
static constexpr size_t CTRL_BUF_SIZE = 1024 * 1024 * 1024; // 1024 MiB
void* ctrl_buf = nullptr;
// RDMA memory region for `gdr_buffer`. Must be nullptr when IBGDA init
// fails.
@ -86,17 +85,6 @@ struct MooncakeEpBuffer {
bool is_roce_ = false;
bool ibgda_disabled_ = false;
int gid_index_ = -1; // Dynamically discovered GID index
int USE_QP_COUNT = MAX_QP_COUNT;
mlx5dv_devx_umem* ctrl_buf_umem;
ibv_pd* pd;
mlx5dv_pd mpd;
memheap* ctrl_buf_heap;
// Fabric memory (MNNVL)
bool use_fabric_mem_ = false;
CUmemGenericAllocationHandle fabric_mem_handle_{};
size_t fabric_alloc_size_ = 0;
// NVLink P2P
int32_t* nvlink_available = nullptr;
@ -168,20 +156,16 @@ struct MooncakeEpBuffer {
return p2p_ipc_all_enabled_;
}
void update_local_qpns();
void sync_ib(const std::vector<int64_t>& remote_addrs,
const std::vector<int32_t>& remote_keys,
const std::vector<int32_t>& remote_qpns,
const std::vector<int32_t>& remote_lids,
const std::vector<int>& active_ranks_mask);
const std::vector<int32_t>& remote_lids);
void sync_roce(const std::vector<int64_t>& remote_addrs,
const std::vector<int32_t>& remote_keys,
const std::vector<int32_t>& remote_qpns,
const std::vector<int64_t>& subnet_prefixes,
const std::vector<int64_t>& interface_ids,
const std::vector<int>& active_ranks_mask);
const std::vector<int64_t>& interface_ids);
std::tuple<int64_t, int32_t> get_mr_info() {
return {(int64_t)mr->addr, (int32_t)mr->rkey};
@ -194,7 +178,7 @@ struct MooncakeEpBuffer {
std::vector<int32_t> get_local_qpns() {
std::vector<int32_t> local_qpns;
for (int i = 0; i < USE_QP_COUNT; ++i) {
for (int i = 0; i < MAX_QP_COUNT; ++i) {
local_qpns.push_back((int32_t)qps[i]->qpn);
}
return local_qpns;
@ -202,7 +186,7 @@ struct MooncakeEpBuffer {
std::vector<int32_t> get_local_lids() {
std::vector<int32_t> local_lids;
for (int i = 0; i < USE_QP_COUNT; ++i) {
for (int i = 0; i < MAX_QP_COUNT; ++i) {
local_lids.push_back((int32_t)qps[i]->port_attr.lid);
}
return local_lids;
@ -210,8 +194,7 @@ struct MooncakeEpBuffer {
std::vector<int32_t> get_ipc_handle();
void sync_nvlink_ipc_handles(
const std::vector<std::vector<int32_t>>& remote_handles,
const std::vector<int>& active_ranks_mask);
const std::vector<std::vector<int32_t>>& remote_handles);
};
inline size_t get_ep_buffer_size_hint(int num_max_dispatch_tokens_per_rank,

View File

@ -5,153 +5,68 @@
#include <stdlib.h>
#include <stdint.h>
#include <stdalign.h>
#include <stdbool.h>
#include <errno.h>
#include "os.h"
#define MEMHEAP_MAX_ALLOCATIONS 1024
struct memheap_allocation {
size_t offset;
size_t size;
bool used;
};
struct memheap {
size_t size;
pthread_mutex_t lock;
size_t allocated;
struct memheap_allocation allocs[MEMHEAP_MAX_ALLOCATIONS];
int alloc_count;
};
static inline struct memheap* memheap_create(size_t size) {
struct memheap* heap = (struct memheap*)malloc(sizeof(struct memheap));
static inline struct memheap *memheap_create(size_t size) {
struct memheap *heap = (struct memheap *)malloc(sizeof(struct memheap));
if (!heap) {
return NULL;
}
heap->size = size;
heap->allocated = 0;
heap->alloc_count = 0;
mutex_init(&heap->lock);
return heap;
}
static inline void memheap_destroy(struct memheap* heap) {
static inline void memheap_destroy(struct memheap *heap) {
if (heap) {
mutex_destroy(&heap->lock);
free(heap);
}
}
static inline size_t memheap_aligned_alloc(struct memheap* heap, size_t size,
static inline size_t memheap_aligned_alloc(struct memheap *heap, size_t size,
size_t align) {
if (size == 0) {
return (size_t)-1; // No allocation for zero size
return 0; // No allocation for zero size
}
if (align == 0 || (align & (align - 1)) != 0) {
errno = EINVAL; // Invalid alignment
return (size_t)-1;
return -1;
}
size_t ret = -1;
mutex_lock(&heap->lock);
size_t ret = (size_t)-1;
for (int i = 0; i < heap->alloc_count; i++) {
if (!heap->allocs[i].used) {
size_t offset = heap->allocs[i].offset;
size_t block_size = heap->allocs[i].size;
size_t aligned_offset = offset;
if (aligned_offset & (align - 1)) {
aligned_offset = (aligned_offset | (align - 1)) + 1;
}
if (aligned_offset + size <= offset + block_size) {
if (aligned_offset > offset) {
int new_idx = heap->alloc_count;
if (new_idx < MEMHEAP_MAX_ALLOCATIONS) {
heap->allocs[new_idx].offset = offset;
heap->allocs[new_idx].size = aligned_offset - offset;
heap->allocs[new_idx].used = false;
heap->alloc_count++;
}
}
if (aligned_offset + size < offset + block_size) {
int new_idx = heap->alloc_count;
if (new_idx < MEMHEAP_MAX_ALLOCATIONS) {
heap->allocs[new_idx].offset = aligned_offset + size;
heap->allocs[new_idx].size =
offset + block_size - (aligned_offset + size);
heap->allocs[new_idx].used = false;
heap->alloc_count++;
}
}
heap->allocs[i].offset = aligned_offset;
heap->allocs[i].size = size;
heap->allocs[i].used = true;
ret = aligned_offset;
heap->allocated += size;
break;
}
}
size_t offset = heap->allocated;
if (offset & (align - 1)) {
offset = (offset | (align - 1)) + 1;
}
if (ret == (size_t)-1) {
size_t offset = heap->allocated;
if (offset & (align - 1)) {
offset = (offset | (align - 1)) + 1;
}
if (offset + size <= heap->size) {
ret = offset;
if (heap->alloc_count < MEMHEAP_MAX_ALLOCATIONS) {
heap->allocs[heap->alloc_count].offset = offset;
heap->allocs[heap->alloc_count].size = size;
heap->allocs[heap->alloc_count].used = true;
heap->alloc_count++;
}
heap->allocated = offset + size;
} else {
errno = ENOMEM;
}
if (offset + size <= heap->size) {
ret = offset;
heap->allocated = offset + size;
} else {
errno = ENOMEM; // Not enough memory
}
mutex_unlock(&heap->lock);
return ret;
}
static inline size_t memheap_alloc(struct memheap* heap, size_t size) {
static inline size_t memheap_alloc(struct memheap *heap, size_t size) {
size_t align = size & -size;
if (align > alignof(max_align_t)) {
align = alignof(max_align_t);
}
if (align < 8) align = 8;
return memheap_aligned_alloc(heap, size, align);
}
static inline void memheap_free(struct memheap* heap, size_t offset) {
if (!heap || offset == (size_t)-1) {
return;
}
mutex_lock(&heap->lock);
for (int i = 0; i < heap->alloc_count; i++) {
if (heap->allocs[i].used && heap->allocs[i].offset == offset) {
heap->allocs[i].used = false;
heap->allocated -= heap->allocs[i].size;
break;
}
}
mutex_unlock(&heap->lock);
static inline void memheap_free(struct memheap *heap, size_t offset) {
// currently no-op
}
#endif
#endif

View File

@ -87,7 +87,7 @@ struct mlx5gda_qp *mlx5gda_create_rc_qp(struct mlx5dv_pd mpd, void *ctrl_buf,
struct memheap *ctrl_buf_heap,
struct ibv_pd *pd, int wqe,
uint8_t port_num, cudaStream_t stream);
void mlx5gda_destroy_qp(struct memheap *ctrl_buf_heap, struct mlx5gda_qp *qp);
void mlx5gda_destroy_qp(struct mlx5gda_qp *qp);
int mlx5gda_modify_rc_qp_rst2init(struct mlx5gda_qp *qp, uint16_t pkey_index);
int mlx5gda_modify_rc_qp_init2rtr(struct mlx5gda_qp *qp,

View File

@ -3,7 +3,7 @@ import re
from setuptools import setup
import torch
from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CUDA_HOME
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
torch_version = re.match(r"\d+(?:\.\d+)*", torch.__version__).group()
@ -13,18 +13,6 @@ module_name = "mooncake.ep" + version_suffix
abi_flag = int(torch._C._GLIBCXX_USE_CXX11_ABI)
current_dir = os.path.abspath(os.path.dirname(__file__))
# Try to link against the CUDA driver stub library if it exists.
cuda_libraries = ["ibverbs", "mlx5"]
cuda_library_dirs = []
if CUDA_HOME is not None:
cuda_stub_dir = os.path.join(CUDA_HOME, "lib64", "stubs")
cuda_stub_lib = os.path.join(cuda_stub_dir, "libcuda.so")
if os.path.exists(cuda_stub_lib):
cuda_libraries.insert(0, "cuda")
cuda_library_dirs.append(cuda_stub_dir)
setup(
name=module_name,
@ -36,7 +24,7 @@ setup(
os.path.join(current_dir, "../mooncake-transfer-engine/include"),
],
sources=[
"src/ep_py.cpp",
"../mooncake-integration/ep/ep_py.cpp",
"src/mooncake_ep_buffer.cpp",
"src/mooncake_ep_kernel.cu",
"src/mooncake_ibgda/mlx5gda.cpp",
@ -45,8 +33,7 @@ setup(
"cxx": [f"-D_GLIBCXX_USE_CXX11_ABI={abi_flag}", "-std=c++20", "-O3", "-g0"],
"nvcc": [f"-D_GLIBCXX_USE_CXX11_ABI={abi_flag}", "-std=c++20", "-Xcompiler", "-O3", "-Xcompiler", "-g0"],
},
libraries=cuda_libraries,
library_dirs=cuda_library_dirs,
libraries=["ibverbs", "mlx5"],
extra_link_args=[
"-Wl,-rpath,$ORIGIN",
"-L" + os.path.join(current_dir, "../mooncake-wheel/mooncake"),

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