forked from mooncake-track/Mooncake
Compare commits
12 Commits
main
...
dev/kv-ind
| Author | SHA1 | Date |
|---|---|---|
|
|
a6bfb2e195 | |
|
|
c837624931 | |
|
|
c67649c83a | |
|
|
6ab4e40ee5 | |
|
|
80c8b41cad | |
|
|
770031b101 | |
|
|
fb7313b87d | |
|
|
c15a911ae0 | |
|
|
45402ff9a1 | |
|
|
fbff581166 | |
|
|
0cc4fea731 | |
|
|
d0ec8173d3 |
|
|
@ -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
|
||||
```
|
||||
|
|
@ -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"
|
||||
|
|
@ -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.
|
||||
|
|
@ -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 \
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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) }}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -5,7 +5,6 @@ build_ofed4
|
|||
old
|
||||
local_test
|
||||
go.sum
|
||||
!mooncake-common/etcd/go.sum
|
||||
*.so
|
||||
bin
|
||||
mod
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 = [
|
||||
|
|
|
|||
121
CMakeLists.txt
121
CMakeLists.txt
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
19
README.md
19
README.md
|
|
@ -15,8 +15,6 @@
|
|||
[](https://kvcache-ai.github.io/Mooncake/)
|
||||
[](https://pypi.org/project/mooncake-transfer-engine)
|
||||
[](https://pypi.org/project/mooncake-transfer-engine)
|
||||
[](https://pypi.org/project/mooncake-transfer-engine)
|
||||
[](https://pypi.org/project/mooncake-transfer-engine-cuda13)
|
||||
[](https://pypi.org/project/mooncake-transfer-engine)
|
||||
[](https://deepwiki.com/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>
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
288
dependencies.sh
288
dependencies.sh
|
|
@ -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}"
|
||||
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
:::
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
@ -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>
|
||||
|
|
@ -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).
|
||||
|
|
@ -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
|
||||
:::
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
@ -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 initiator’s launch command**:
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
@ -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+
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
@ -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 |
|
|
@ -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
|
||||
:::
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||

|
||||
|
||||
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.
|
||||
|
||||

|
||||
|
||||
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`.
|
||||
|
|
@ -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+
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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"。
|
||||
|
||||
|
|
|
|||
|
|
@ -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 组件
|
||||
|
|
|
|||
|
|
@ -1,90 +0,0 @@
|
|||
# Kunpeng UB Transport
|
||||
Kunpeng UbTransport源代码路径为Mooncake/mooncake-transfer-engine/src/transport/kunpneg_transport,该路径下有UB协议的Transport对接代码和实现逻辑。
|
||||
|
||||
## 概述
|
||||
UB(Unified Bus,统一总线) 是与RDMA、CXL、NVLink 和TCP处于同一抽象层的传输协议,属于可在应用层灵活选择的传输方案。目前 UB 协议有两个开源实现:URMA(远程内存访问语义)和 OBMM(Load/Store 语义)。
|
||||
|
||||
URMA(Unified 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
|
||||
```
|
||||
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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}")
|
||||
|
|
@ -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()
|
||||
|
|
@ -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}")
|
||||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
)
|
||||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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=
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
)
|
||||
|
|
@ -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()
|
||||
}
|
||||
|
|
@ -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
|
||||
)
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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() {}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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()
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
}
|
||||
}
|
||||
|
|
@ -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"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
}'
|
||||
|
||||
```
|
||||
|
|
@ -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
|
||||
|
|
@ -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())
|
||||
}
|
||||
|
|
@ -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())
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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")
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
)
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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()
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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")
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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")
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
Loading…
Reference in New Issue