diff --git a/.ci/run_test.sh b/.ci/run_test.sh deleted file mode 100644 index afeac3a31a..0000000000 --- a/.ci/run_test.sh +++ /dev/null @@ -1,51 +0,0 @@ -set -ex -TEST_DIR=$1 -OUTPUT_DIR=$2 - - -run_test(){ - echo "$OUTPUT_DIR" - ./"$1" - LOCAL_SUCCESS=$? - if [[ $LOCAL_SUCCESS == 0 ]]; then - echo "PASSED : $1 " >> "$OUTPUT_DIR"/results.txt - else - echo "FAILED : $1 " >> "$OUTPUT_DIR"/results.txt - fi -} - - -pushd $TEST_DIR -etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://10.0.0.1:2379 & -export MC_GID_INDEX=1 -run_test transport_uint_test -run_test tcp_transport_test -export MC_GID_INDEX=1 - -# sudo pkill etcd -# sleep 1 -# etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://10.0.0.1:2379 & -# sudo fuser -k 14345/tcp -export MC_GID_INDEX=1 && ./rdma_transport_test --mode=target --metadata_server=127.0.0.1:2379 --local_server_name=127.0.0.2:14345 --device_name=erdma_0 & -export MC_GID_INDEX=1 && ./rdma_transport_test --metadata_server=127.0.0.1:2379 --segment_id=127.0.0.2:14345 --local_server_name=127.0.0.3:14346 --device_name=erdma_1 -LOCAL_SUCCESS=$? -if [[ $LOCAL_SUCCESS == 0 ]]; then - echo "PASSED : rdma_transport_test " >> "$OUTPUT_DIR"/results.txt -else - echo "FAILED : rdma_transport_test " >> "$OUTPUT_DIR"/results.txt -fi - -# sudo pkill etcd -# sleep 1 -# etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://10.0.0.1:2379 & -# export MC_GID_INDEX=1 -# sudo fuser -k 15345/tcp -# export MC_GID_INDEX=1 && ./rdma_transport_test --mode=target --metadata_server=127.0.0.1:2379 --local_server_name=127.0.0.2:14345 --device_name=erdma_0 & -export MC_GID_INDEX=1 && ./rdma_transport_test2 --metadata_server=127.0.0.1:2379 --segment_id=127.0.0.2:14345 --local_server_name=127.0.0.3:14346 --device_name=erdma_1 -LOCAL_SUCCESS=$? -if [[ $LOCAL_SUCCESS == 0 ]]; then - echo "PASSED : rdma_transport_test2 " >> "$OUTPUT_DIR"/results.txt -else - echo "FAILED : rdma_transport_test2 " >> "$OUTPUT_DIR"/results.txt -fi -popd \ No newline at end of file diff --git a/.claude-plugin/marketplace.json b/.claude-plugin/marketplace.json new file mode 100644 index 0000000000..e15c855059 --- /dev/null +++ b/.claude-plugin/marketplace.json @@ -0,0 +1,42 @@ +{ + "$schema": "https://json.schemastore.org/claude-code-marketplace.json", + "name": "mooncake", + "owner": { + "name": "Mooncake Team" + }, + "metadata": { + "description": "Built-in Claude Code skills for working with Mooncake: deployment troubleshooting, local CI validation, and the Python API." + }, + "plugins": [ + { + "name": "mooncake-troubleshoot", + "source": { + "source": "git-subdir", + "url": "https://github.com/kvcache-ai/Mooncake.git", + "path": ".claude/skills/mooncake-troubleshoot" + }, + "description": "Systematically diagnose Mooncake deployment and runtime issues (services, RDMA, env vars, connectivity, logs) and propose fixes.", + "homepage": "https://kvcache-ai.github.io/Mooncake/skills/mooncake-troubleshoot.html" + }, + { + "name": "mooncake-ci-local", + "source": { + "source": "git-subdir", + "url": "https://github.com/kvcache-ai/Mooncake.git", + "path": ".claude/skills/mooncake-ci-local" + }, + "description": "Run Mooncake pre-PR local validation via scripts/run_ci_test.sh, reproducing the reproducible parts of GitHub Actions.", + "homepage": "https://kvcache-ai.github.io/Mooncake/skills/mooncake-ci-local.html" + }, + { + "name": "mooncake-api", + "source": { + "source": "git-subdir", + "url": "https://github.com/kvcache-ai/Mooncake.git", + "path": ".claude/skills/mooncake-api" + }, + "description": "Work with the Mooncake Store, Transfer Engine, and EP/Backend Python APIs.", + "homepage": "https://kvcache-ai.github.io/Mooncake/skills/mooncake-api.html" + } + ] +} diff --git a/.claude/skills/mooncake-api/.claude-plugin/plugin.json b/.claude/skills/mooncake-api/.claude-plugin/plugin.json new file mode 100644 index 0000000000..81bad70112 --- /dev/null +++ b/.claude/skills/mooncake-api/.claude-plugin/plugin.json @@ -0,0 +1,12 @@ +{ + "$schema": "https://json.schemastore.org/claude-code-plugin-manifest.json", + "name": "mooncake-api", + "description": "Work with the Mooncake Store, Transfer Engine, and EP/Backend Python APIs.", + "author": { + "name": "Mooncake Team" + }, + "homepage": "https://kvcache-ai.github.io/Mooncake/skills/mooncake-api.html", + "repository": "https://github.com/kvcache-ai/Mooncake", + "license": "Apache-2.0", + "keywords": ["mooncake", "python", "api", "transfer-engine", "store"] +} diff --git a/.claude/skills/mooncake-api/SKILL.md b/.claude/skills/mooncake-api/SKILL.md index daa1b6c51b..f59b189623 100644 --- a/.claude/skills/mooncake-api/SKILL.md +++ b/.claude/skills/mooncake-api/SKILL.md @@ -1,4 +1,7 @@ - +--- +name: mooncake-api +description: Help users work with the Mooncake Python APIs for distributed storage and high-performance data transfer. Use when working with Mooncake Store (distributed KV cache), Transfer Engine (RDMA/TCP transfers), service setup (master, metadata server), PyTorch tensors in the Store, zero-copy/buffer management, batch operations and replication, or Mooncake EP / Backend (Expert Parallelism). Trigger on questions about MooncakeDistributedStore, TransferEngine, put/get, put_tensor, register_buffer, ReplicateConfig, or the mooncake.store / mooncake.engine / mooncake.pg Python modules. +--- # Mooncake Python API Skill @@ -16,6 +19,14 @@ Use this skill when users ask about: - Mooncake EP (Expert Parallelism) and Mooncake Backend - Troubleshooting Mooncake Python API issues +## Routing Guidance + +- For most vLLM and SGLang users, start from `docs/source/getting_started/quick-start.md`. +- For PD disaggregation, direct users to the SGLang/vLLM integration guides listed in Quick Start. Those guides own the serving-framework configuration. +- For Mooncake Store integrations, direct users to the SGLang/vLLM Store setup guides listed in Quick Start. Do not duplicate `mooncake_master` startup commands in general API answers unless the user is working outside those frameworks. +- For direct low-level Transfer Engine usage, use `docs/source/design/transfer-engine/index.md#using-transfer-engine-in-your-projects`. +- For API signatures and method details, use the Python API references under `docs/source/python-api-reference/`. + ## Core Components ### 1. Mooncake Store (Distributed KV Cache) @@ -304,6 +315,7 @@ mooncake_master --default_kv_lease_ttl=5000 - `MC_STORE_MEMCPY`: Enable local memcpy optimization (set to "1") - `MC_STORE_CLIENT_METRIC`: Enable client metrics (enabled by default) - `MC_YLT_LOG_LEVEL`: Log level (trace/debug/info/warn/error/critical) +- `MOONCAKE_STORE_CHECKSUM`: Enable diagnostic object-level CRC-64 checks (set to "1" before creating any writer or reader client) ## Common Patterns @@ -462,6 +474,17 @@ if result != 0: raise RuntimeError(f"Failed to register buffer: {result}") ``` +### Corrupted Data or Garbled Output +```python +# Set this before importing Mooncake or creating any Store client. +import os +os.environ["MOONCAKE_STORE_CHECKSUM"] = "1" + +from mooncake.store import MooncakeDistributedStore +``` + +Enable the switch on every writer and reader client process, then reproduce with full-object `put`/`upsert` and `get` operations. Treat `CHECKSUM_MISMATCH` (-801) as a failed read and do not use the destination buffer. Objects without checksum metadata and range reads are not verified. This mode scans object data, stages GPU buffers to host memory, and disables the local hot cache, so use it only for diagnosis. + ### Service Connectivity ```bash # Check master is running @@ -507,6 +530,8 @@ curl http://localhost:8080/metadata ## Documentation Links - Full API Reference: https://kvcache-ai.github.io/Mooncake/ +- Quick Start: docs/source/getting_started/quick-start.md - Mooncake Store: docs/source/python-api-reference/mooncake-store.md - Transfer Engine: docs/source/python-api-reference/transfer-engine.md +- Transfer Engine direct usage: docs/source/design/transfer-engine/index.md#using-transfer-engine-in-your-projects - EP Backend: docs/source/python-api-reference/ep-backend.md diff --git a/.claude/skills/mooncake-ci-local/.claude-plugin/plugin.json b/.claude/skills/mooncake-ci-local/.claude-plugin/plugin.json new file mode 100644 index 0000000000..afb3699b20 --- /dev/null +++ b/.claude/skills/mooncake-ci-local/.claude-plugin/plugin.json @@ -0,0 +1,12 @@ +{ + "$schema": "https://json.schemastore.org/claude-code-plugin-manifest.json", + "name": "mooncake-ci-local", + "description": "Run Mooncake pre-PR local validation via scripts/run_ci_test.sh, reproducing the reproducible parts of GitHub Actions.", + "author": { + "name": "Mooncake Team" + }, + "homepage": "https://kvcache-ai.github.io/Mooncake/skills/mooncake-ci-local.html", + "repository": "https://github.com/kvcache-ai/Mooncake", + "license": "Apache-2.0", + "keywords": ["mooncake", "ci", "validation", "pre-pr"] +} diff --git a/.claude/skills/mooncake-ci-local/SKILL.md b/.claude/skills/mooncake-ci-local/SKILL.md new file mode 100644 index 0000000000..546bcc24af --- /dev/null +++ b/.claude/skills/mooncake-ci-local/SKILL.md @@ -0,0 +1,159 @@ +--- +name: mooncake-ci-local +description: Run Mooncake pre-PR local validation through scripts/run_ci_test.sh. Use this skill whenever the user wants to validate a branch before opening or submitting a PR, run local CI, run ci test, check changes before PR, reproduce GitHub Actions locally, or force a full pre-submit verification. Trigger on phrases like "提交 PR 前验证", "run ci test", "run local CI", "check my branch", "test before PR", "pre-submit validation", and "reproduce CI locally". +--- + +# Mooncake Pre-PR Local Validation + +Use `bash scripts/run_ci_test.sh` as the default entry point. This is the single local lane for PR-before-submit validation, and it already coordinates the reproducible parts of GitHub Actions. + +## Default Entry Point + +When the user asks for any of the following, run the repo script first instead of reconstructing the workflow by hand: + +- 提交 PR 前本地验证 +- run ci test +- run local CI +- check my branch before PR +- reproduce CI locally + +Default command: + +```bash +bash scripts/run_ci_test.sh +``` + +What this script already covers: + +- GitHub-like `paths-filter` against `origin/main` +- `typos` +- `scripts/code_format.sh --check` +- default CMake configure/build/install in `build-ci-local` +- `ctest` +- wheel build in `build-wheel-local` +- wheel installation validation +- `scripts/run_tests.sh` +- selected Python API and integration tests +- per-stage summary and logs under `local_test/run-ci-logs//` + +## Standard Agent Workflow + +1. Run `bash scripts/run_ci_test.sh` from the repo root unless the user explicitly asks for a narrower subset. +2. Read the stage summary instead of dumping raw terminal output. +3. Report these items back to the user: + - passed stages + - failed stages + - blocked stages + - unsupported stages + - whether `paths-filter` skipped downstream stages + - the log directory under `local_test/run-ci-logs/...` +4. If there is a failure, inspect the corresponding stage log and summarize the root cause. + +## Common Options + +Force a full lane even if `paths-filter` would skip downstream stages: + +```bash +bash scripts/run_ci_test.sh --skip-path-filter +``` + +Use another base ref: + +```bash +bash scripts/run_ci_test.sh --base origin/main +``` + +Auto-install missing dependencies: + +```bash +bash scripts/run_ci_test.sh --install-deps +``` + +Keep services running for follow-up debugging: + +```bash +bash scripts/run_ci_test.sh --keep-services +``` + +## Minimal Example + +User prompt: + +- 提交 PR 前，帮我跑一遍本地 CI 验证当前分支。 + +Expected action: + +```bash +bash scripts/run_ci_test.sh +``` + +If the user wants to ignore changed-path optimization and force the full lane: + +```bash +bash scripts/run_ci_test.sh --skip-path-filter +``` + +See also `.claude/skills/mooncake-ci-local/examples/minimal.md`. + +## How To Interpret Results + +- `passed`: the stage succeeded locally. +- `failed`: the stage reproduced a real local failure and needs investigation. +- `blocked`: local environment or dependency issue prevented execution. +- `unsupported`: intentionally not run in the local lane because it needs external platforms, special hardware, or a non-default build. + +If `paths-filter` skips downstream stages, explain that the current branch changed only non-source paths relative to the selected base. + +## Current Local Coverage + +Included by default: + +- spell check +- code format check +- default ASan CMake lane in `build-ci-local` +- `ctest` +- wheel build and installation test +- `scripts/run_tests.sh` +- selected Python API tests + +Unsupported by design in the default local lane: + +- Ascend jobs +- T-one integration jobs +- MUSA jobs +- Docker image build jobs +- CUDA 13 wheel jobs +- PG-backend tests absent from the default wheel build +- Python drain-http API stage in the local ASan lane + +## Targeted Reruns For Debugging + +Use targeted reruns only after the full script identifies a failing area, or when the user explicitly asks for a smaller scope. + +Rerun a specific C++ test pattern: + +```bash +cd build-ci-local +MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 ctest -R --output-on-failure +``` + +Rerun the Python wheel integration lane: + +```bash +source test_env/bin/activate +MC_STORE_MEMCPY=false TEST_SSD_OFFLOAD_IN_EVICT=true ./scripts/run_tests.sh +``` + +Rerun the safetensor unittest: + +```bash +source test_env/bin/activate +python -m unittest mooncake-wheel.tests.test_safetensor_functions +``` + +## Notes For The Agent + +- Prefer the repo script over rebuilding the CI workflow step by step. +- Preserve the separation between `build-ci-local` and `build-wheel-local`. +- Summarize failing stages from their logs instead of pasting raw output. +- If the user only asks whether the branch is safe before opening a PR, the default answer path is `bash scripts/run_ci_test.sh`. diff --git a/.claude/skills/mooncake-ci-local/examples/minimal.md b/.claude/skills/mooncake-ci-local/examples/minimal.md new file mode 100644 index 0000000000..3d081c736b --- /dev/null +++ b/.claude/skills/mooncake-ci-local/examples/minimal.md @@ -0,0 +1,30 @@ +# Minimal Example + +Goal: validate the current branch before opening or submitting a PR. + +User prompt: + +- 提交 PR 前，帮我跑一遍本地 CI 验证当前分支。 + +Expected action: + +```bash +bash scripts/run_ci_test.sh +``` + +Force a full rerun even if `paths-filter` would skip downstream stages: + +```bash +bash scripts/run_ci_test.sh --skip-path-filter +``` + +Result location: + +- `local_test/run-ci-logs//` + +Typical report format: + +- passed stages +- failed or blocked stages +- unsupported stages +- first actionable failure and its log path diff --git a/.claude/skills/mooncake-ci-local/scripts/check-prerequisites.sh b/.claude/skills/mooncake-ci-local/scripts/check-prerequisites.sh new file mode 100644 index 0000000000..5f318e43aa --- /dev/null +++ b/.claude/skills/mooncake-ci-local/scripts/check-prerequisites.sh @@ -0,0 +1,101 @@ +#!/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" diff --git a/.claude/skills/mooncake-troubleshoot/.claude-plugin/plugin.json b/.claude/skills/mooncake-troubleshoot/.claude-plugin/plugin.json new file mode 100644 index 0000000000..239659ad8c --- /dev/null +++ b/.claude/skills/mooncake-troubleshoot/.claude-plugin/plugin.json @@ -0,0 +1,12 @@ +{ + "$schema": "https://json.schemastore.org/claude-code-plugin-manifest.json", + "name": "mooncake-troubleshoot", + "description": "Systematically diagnose Mooncake deployment and runtime issues (services, RDMA, env vars, connectivity, logs) and propose fixes.", + "author": { + "name": "Mooncake Team" + }, + "homepage": "https://kvcache-ai.github.io/Mooncake/skills/mooncake-troubleshoot.html", + "repository": "https://github.com/kvcache-ai/Mooncake", + "license": "Apache-2.0", + "keywords": ["mooncake", "troubleshooting", "rdma", "deployment"] +} diff --git a/.claude/skills/mooncake-troubleshoot/SKILL.md b/.claude/skills/mooncake-troubleshoot/SKILL.md new file mode 100644 index 0000000000..c95410f035 --- /dev/null +++ b/.claude/skills/mooncake-troubleshoot/SKILL.md @@ -0,0 +1,383 @@ +--- +name: mooncake-troubleshoot +description: Automatically diagnose Mooncake deployment and runtime issues. Checks services (mooncake_master, metadata server), RDMA devices, environment variables, connectivity, memory limits, object integrity, and analyzes logs for common error patterns. Use when Mooncake deployment fails, services won't start, connections fail, data is corrupted or garbled, or you encounter runtime errors like "Error from etcd client", "No matched device found", "Failed to register memory", "NO_AVAILABLE_HANDLE", "CHECKSUM_MISMATCH", 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://:2379 + ``` + - **Fix:** Disable HTTP proxy: + ```bash + unset http_proxy https_proxy + ``` + +### 3. Environment Variables Check + +Verify critical environment variables are set correctly: + +```bash +# Display Mooncake variables +env | grep -E '^(MC_|MOONCAKE_STORE_CHECKSUM=)' + +# 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" +echo "MOONCAKE_STORE_CHECKSUM: $MOONCAKE_STORE_CHECKSUM" +``` + +**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") +- `MOONCAKE_STORE_CHECKSUM=1` - Enable diagnostic object-level CRC-64 checks; set before starting every writer and reader client process + +### 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 -R + +# On peer node: +ib_write_bw -d -R + +# Test GPU Direct RDMA (if CUDA enabled) +ib_write_bw -d -R -x gdr + +# On peer node: +ib_write_bw -d -R -x gdr + +# Test DNS resolution +nslookup +ping +``` + +**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 +- `CHECKSUM_MISMATCH` (-801) → Full-object read differs from the checksum captured before the write + - **Action:** Treat the read as failed and do not use the destination buffer +- `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 + +### 9. Object Integrity Diagnostics + +When the user reports corrupted data or garbled output that may originate from Mooncake Store: + +1. Verify that the Store clients, primary master, and standby master use checksum-capable binaries from the same version. +2. Set `MOONCAKE_STORE_CHECKSUM=1` before starting every writer and reader client process, then restart or recreate existing clients. +3. Reproduce with full-object `put`/`upsert` and `get` operations. Range reads, including `get_into_ranges`, are not covered. +4. If `CHECKSUM_MISMATCH` (-801) is returned, the covered read differs from the checksum captured over the writer's source object. Treat the read as failed and do not consume the destination buffer. +5. If `object_checksum_absent` appears at VLOG(1), metadata has no checksum and verification was skipped. This can occur for existing objects or objects written by clients without the switch. +6. If verification succeeds, continue investigating corruption outside the covered Store path. A checksum mismatch identifies a difference but does not by itself distinguish transfer corruption from storage corruption. +7. Warn that this diagnostic mode scans all object data, stages GPU buffers to host memory, and disables the local hot cache. Disable it after diagnosis. + +## 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 | +| -801 | CHECKSUM_MISMATCH | Full-object data differs from its stored checksum | Reject the read buffer and investigate the Store data path | +| -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://: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. diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile index 51cf93e108..5078333c83 100644 --- a/.devcontainer/Dockerfile +++ b/.devcontainer/Dockerfile @@ -8,11 +8,10 @@ RUN apt-get update -y \ RUN apt-get install -y libibverbs-dev \ libunwind-dev \ libgoogle-glog-dev \ - libgtest-dev \ libjsoncpp-dev \ libnuma-dev \ libpython3-dev \ - libboost-all-dev \ + libboost-dev \ libssl-dev \ libgrpc-dev \ libgrpc++-dev \ @@ -23,19 +22,21 @@ RUN apt-get install -y libibverbs-dev \ libhiredis-dev \ libyaml-cpp-dev \ libjemalloc-dev \ + libzstd-dev \ + libmsgpack-dev \ + libgflags-dev \ + libxxhash-dev \ pkg-config \ patchelf -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 \ - && mkdir -p build \ - && cd build \ - && cmake .. -DBUILD_EXAMPLES=OFF -DBUILD_BENCHMARK=OFF -DBUILD_UNIT_TESTS=OFF \ - && cmake --build . -j$(nproc) \ - && cmake --install . +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 ENV GOPROXY='https://goproxy.cn' ENV PATH=/usr/local/go/bin:$PATH diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json index 7df63069f3..dea0cbb24b 100644 --- a/.devcontainer/devcontainer.json +++ b/.devcontainer/devcontainer.json @@ -1,6 +1,8 @@ { "name": "Mooncake Dev", - "dockerFile": "Dockerfile", + "build": { + "dockerfile": "Dockerfile" + }, "runArgs": [ "--cap-add=SYS_PTRACE", "--cap-add=NET_RAW", @@ -18,7 +20,7 @@ "vscode": { "settings": { "terminal.integrated.shell.linux": "/bin/bash", - "bazel.buildifierFixOnFormat": true, + "bazel.buildifierFixOnFormat": true } } } diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index f6c5f05b1e..3d83b130e1 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -6,16 +6,24 @@ # EP: @UNIDY2002 UNIDY2002@outlook.com # PG: @UNIDY2002 UNIDY2002@outlook.com -.github @stmatengss @ykwd @Ann-1024 @luketong777 -/docs @ShangmingCai @stmatengss @ykwd -/mooncake-ep @UNIDY2002 @ympcMark -/mooncake-integration/ep @UNIDY2002 @ympcMark -/mooncake-integration/transfer_engine @ShangmingCai @alogfans -/mooncake-integration/store @ykwd @stmatengss -/mooncake-pg @UNIDY2002 @ympcMark +.github @stmatengss @ykwd @Ann-1024 @luketong777 @Aionw +/docs @ShangmingCai @stmatengss @ykwd @UNIDY2002 @alogfans @staryxchen +/mooncake-ep @UNIDY2002 @ympcMark @yuechen-sys +/mooncake-integration/transfer_engine @ShangmingCai @alogfans +/mooncake-integration/store @ykwd @stmatengss @zxpdemonio +/mooncake-pg @UNIDY2002 @ympcMark @yuechen-sys +/mooncake-reshard @ShangmingCai @stmatengss @Bo-Vincent @zxpdemonio /mooncake-store @ykwd @stmatengss @XucSh @YiXR -/mooncake-transfer-engine @alogfans @doujiang24 @chestnut-Q +/mooncake-store/*/ha/ @Libotry @YiXR @00fish0 @Icedcoco @Aionw +/python/mooncake/ep.py @UNIDY2002 @ympcMark @yuechen-sys +/python/mooncake/mooncake_elastic_buffer.py @UNIDY2002 @ympcMark @yuechen-sys +/python/mooncake/mooncake_ep_buffer.py @UNIDY2002 @ympcMark @yuechen-sys +/python/tests/ep/ @UNIDY2002 @ympcMark @yuechen-sys +/mooncake-transfer-engine @alogfans @doujiang24 @chestnut-Q @staryxchen +/mooncake-transfer-engine/tent @alogfans @doujiang24 @chestnut-Q @staryxchen @00fish0 @dtcccc /mooncake-transfer-engine/*/transport/hip_transport/ @alogfans @amd-arozanov /mooncake-transfer-engine/*/transport/ascend_transport/ @alogfans @ascend-direct-dev -/mooncake-wheel @ShangmingCai @stmatengss +/mooncake-transfer-engine/*/transport/efa_transport/ @alogfans @whn09 +/mooncake-wheel @ShangmingCai @stmatengss @zxpdemonio /scripts/tone_tests @luketong777 +/scripts/ascend/ @ascend-direct-dev @VNightMare @MingYang119 diff --git a/.github/actions/ctest-diagnostics/action.yml b/.github/actions/ctest-diagnostics/action.yml new file mode 100644 index 0000000000..eca7e3d6ae --- /dev/null +++ b/.github/actions/ctest-diagnostics/action.yml @@ -0,0 +1,50 @@ +name: Preserve CTest failure diagnostics +description: Summarize a failed CTest run and upload its logs and JUnit report + +inputs: + summary-title: + description: Heading to add to the GitHub Actions job summary + required: true + artifact-name: + description: Name of the uploaded diagnostics artifact + required: true + junit-report: + description: Path to the CTest JUnit XML report + required: true + failed-tests-log: + description: Path to CTest's LastTestsFailed.log + required: true + last-test-log: + description: Path to CTest's LastTest.log + required: true + retention-days: + description: Number of days to retain the diagnostics artifact + required: false + default: "14" + +runs: + using: composite + steps: + - name: Add CTest failure summary + env: + SUMMARY_TITLE: ${{ inputs.summary-title }} + FAILED_TESTS_LOG: ${{ inputs.failed-tests-log }} + LAST_TEST_LOG: ${{ inputs.last-test-log }} + shell: bash + run: | + "$GITHUB_ACTION_PATH/summarize.sh" \ + "$SUMMARY_TITLE" \ + "$FAILED_TESTS_LOG" \ + "$LAST_TEST_LOG" + + - name: Upload CTest diagnostics + if: ${{ always() }} + uses: actions/upload-artifact@v4 + with: + name: ${{ inputs.artifact-name }} + if-no-files-found: warn + retention-days: ${{ inputs.retention-days }} + path: | + ${{ inputs.junit-report }} + ${{ inputs.last-test-log }} + ${{ inputs.failed-tests-log }} diff --git a/.github/actions/ctest-diagnostics/summarize.sh b/.github/actions/ctest-diagnostics/summarize.sh new file mode 100755 index 0000000000..d1d0c4b1dc --- /dev/null +++ b/.github/actions/ctest-diagnostics/summarize.sh @@ -0,0 +1,58 @@ +#!/usr/bin/env bash + +set -euo pipefail + +if [[ $# -ne 3 ]]; then + echo "Usage: $0 " >&2 + exit 2 +fi + +: "${GITHUB_STEP_SUMMARY:?GITHUB_STEP_SUMMARY must be set}" + +readonly SUMMARY_TITLE="$1" +readonly FAILED_TESTS_LOG="$2" +readonly LAST_TEST_LOG="$3" + +{ + echo "## $SUMMARY_TITLE" + echo + + if [[ -s "$FAILED_TESTS_LOG" ]]; then + echo "### Failed CTest targets" + echo '```text' + cat "$FAILED_TESTS_LOG" + echo '```' + fi + + if [[ -s "$LAST_TEST_LOG" ]]; then + failed_cases=$(grep -E '^\[ FAILED \]' "$LAST_TEST_LOG" | + sed -E \ + -e '/^\[ FAILED \] [0-9]+ tests?, listed below:$/d' \ + -e 's/ \([0-9]+ ms\)$//' | + sort -u || true) + if [[ -n "$failed_cases" ]]; then + echo "### Failed test cases" + echo '```text' + echo "$failed_cases" + echo '```' + fi + + failure_context=$(grep -n -B2 -A8 -E \ + -e ':[0-9]+: Failure$' \ + -e 'Assertion .* failed\.$' \ + -e '(ERROR|WARNING): (AddressSanitizer|LeakSanitizer|ThreadSanitizer|UndefinedBehaviorSanitizer)' \ + -e 'runtime error:|Segmentation fault|terminate called' \ + "$LAST_TEST_LOG" | sed -n '1,200p' || true) + if [[ -n "$failure_context" ]]; then + echo "### Failure context" + echo '```text' + echo "$failure_context" + echo '```' + fi + else + echo "CTest did not produce LastTest.log. Check the failed step for setup errors." + fi + + echo + echo "Download the CTest diagnostics artifact for the complete log and JUnit report." +} >> "$GITHUB_STEP_SUMMARY" diff --git a/.github/actions/ctest-diagnostics/test.sh b/.github/actions/ctest-diagnostics/test.sh new file mode 100755 index 0000000000..745f0fb5f1 --- /dev/null +++ b/.github/actions/ctest-diagnostics/test.sh @@ -0,0 +1,16 @@ +#!/usr/bin/env bash + +set -euo pipefail + +readonly ACTION_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +readonly FIXTURE_DIR="$ACTION_DIR/testdata" +readonly SUMMARY_FILE="$(mktemp)" + +trap 'rm -f "$SUMMARY_FILE"' EXIT + +GITHUB_STEP_SUMMARY="$SUMMARY_FILE" "$ACTION_DIR/summarize.sh" \ + "CTest failure fixture" \ + "$FIXTURE_DIR/last-tests-failed.txt" \ + "$FIXTURE_DIR/last-test.txt" + +diff -u "$FIXTURE_DIR/expected-summary.md" "$SUMMARY_FILE" diff --git a/.github/actions/ctest-diagnostics/testdata/expected-summary.md b/.github/actions/ctest-diagnostics/testdata/expected-summary.md new file mode 100644 index 0000000000..24cb45e6da --- /dev/null +++ b/.github/actions/ctest-diagnostics/testdata/expected-summary.md @@ -0,0 +1,43 @@ +## CTest failure fixture + +### Failed CTest targets +```text +1:google_test +2:native_assertion +3:tsan_test +``` +### Failed test cases +```text +[ FAILED ] FooTest.Bar +``` +### Failure context +```text +3-1/3 Testing: google_test +4-[ RUN ] FooTest.Bar +5:/workspace/foo_test.cpp:42: Failure +6-Expected equality of these values: +7- 1 +8- actual +9- Which is: 2 +10-[ FAILED ] FooTest.Bar (5 ms) +11-[==========] 1 test from 1 test suite ran. (5 ms total) +12-[ PASSED ] 0 tests. +13-[ FAILED ] 1 test, listed below: +-- +15- +16-2/3 Testing: native_assertion +17:native_assertion: /workspace/assert.cpp:17: int main(): Assertion 'ready' failed. +18-Subprocess aborted +19- +20-3/3 Testing: tsan_test +21:WARNING: ThreadSanitizer: data race (pid=123) +22- Write of size 4 at 0x7b0400000800 by thread T1: +23- #0 worker /workspace/worker.cpp:9 +24- +25-SUMMARY: ThreadSanitizer: data race /workspace/worker.cpp:9 in worker +26-ThreadSanitizer: reported 1 warnings +27- +28-End testing +``` + +Download the CTest diagnostics artifact for the complete log and JUnit report. diff --git a/.github/actions/ctest-diagnostics/testdata/last-test.txt b/.github/actions/ctest-diagnostics/testdata/last-test.txt new file mode 100644 index 0000000000..4331103b20 --- /dev/null +++ b/.github/actions/ctest-diagnostics/testdata/last-test.txt @@ -0,0 +1,28 @@ +Start testing: fixture +---------------------------------------------------------- +1/3 Testing: google_test +[ RUN ] FooTest.Bar +/workspace/foo_test.cpp:42: Failure +Expected equality of these values: + 1 + actual + Which is: 2 +[ FAILED ] FooTest.Bar (5 ms) +[==========] 1 test from 1 test suite ran. (5 ms total) +[ PASSED ] 0 tests. +[ FAILED ] 1 test, listed below: +[ FAILED ] FooTest.Bar + +2/3 Testing: native_assertion +native_assertion: /workspace/assert.cpp:17: int main(): Assertion 'ready' failed. +Subprocess aborted + +3/3 Testing: tsan_test +WARNING: ThreadSanitizer: data race (pid=123) + Write of size 4 at 0x7b0400000800 by thread T1: + #0 worker /workspace/worker.cpp:9 + +SUMMARY: ThreadSanitizer: data race /workspace/worker.cpp:9 in worker +ThreadSanitizer: reported 1 warnings + +End testing diff --git a/.github/actions/ctest-diagnostics/testdata/last-tests-failed.txt b/.github/actions/ctest-diagnostics/testdata/last-tests-failed.txt new file mode 100644 index 0000000000..dfb4965678 --- /dev/null +++ b/.github/actions/ctest-diagnostics/testdata/last-tests-failed.txt @@ -0,0 +1,3 @@ +1:google_test +2:native_assertion +3:tsan_test diff --git a/.github/actions/free-disk-space/action.yml b/.github/actions/free-disk-space/action.yml new file mode 100644 index 0000000000..51767484ed --- /dev/null +++ b/.github/actions/free-disk-space/action.yml @@ -0,0 +1,15 @@ +name: Free runner disk space +description: Remove preinstalled toolchains that Mooncake CI does not use + +runs: + using: composite + steps: + - name: Remove unused runner tools + shell: bash + run: | + sudo rm -rf -- \ + /usr/share/dotnet \ + /opt/ghc \ + /opt/hostedtoolcache/CodeQL \ + /usr/local/lib/android + df -h diff --git a/.github/labeler.yml b/.github/labeler.yml index dad607da0d..36fc9db2ee 100644 --- a/.github/labeler.yml +++ b/.github/labeler.yml @@ -17,8 +17,57 @@ PyTorch Backend: Mooncake EP: - changed-files: - - any-glob-to-any-file: 'mooncake-ep/**/*' + - any-glob-to-any-file: + - 'mooncake-ep/**/*' + - 'python/mooncake/ep.py' + - 'python/mooncake/mooncake_elastic_buffer.py' + - 'python/mooncake/mooncake_ep_buffer.py' + - 'python/tests/ep/**/*' Installation: - changed-files: - any-glob-to-any-file: 'mooncake-wheel/**/*' + +P2P Store: + - changed-files: + - any-glob-to-any-file: 'mooncake-p2p-store/**/*' + +Integration: + - changed-files: + - any-glob-to-any-file: 'mooncake-integration/**/*' + +Common: + - changed-files: + - any-glob-to-any-file: + - 'mooncake-common/**/*' + - 'mooncake-reshard/**/*' + +CI/Build: + - changed-files: + - any-glob-to-any-file: + - '.github/**/*' + - 'CMakeLists.txt' + - '.pre-commit-config.yaml' + - 'dependencies.sh' + +Documentation: + - changed-files: + - any-glob-to-any-file: + - 'docs/**/*' + - '*.md' + +Tests: + - changed-files: + - any-glob-to-any-file: + - 'scripts/test_*' + - 'mooncake-wheel/tests/**/*' + - 'mooncake-reshard/tests/**/*' + - 'python/tests/**/*' + - 'scripts/tone_tests/**/*' + +Ascend/NPU: + - changed-files: + - any-glob-to-any-file: + - '**/ascend*/**' + - 'scripts/ascend/**' + - '.github/workflows/*ascend*' diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 541ebfd245..0ca038db46 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -6,11 +6,13 @@ - [ ] Transfer Engine (`mooncake-transfer-engine`) - [ ] Mooncake Store (`mooncake-store`) +- [ ] Reshard (`mooncake-reshard`) - [ ] Mooncake EP (`mooncake-ep`) +- [ ] Mooncake PG (`mooncake-pg`) - [ ] Integration (`mooncake-integration`) - [ ] P2P Store (`mooncake-p2p-store`) - [ ] Python Wheel (`mooncake-wheel`) -- [ ] PyTorch Backend (`mooncake-pg`) +- [ ] Common (`mooncake-common`) - [ ] Mooncake RL (`mooncake-rl`) - [ ] CI/CD - [ ] Docs @@ -23,15 +25,38 @@ - [ ] Refactor - [ ] Breaking change - [ ] Documentation update +- [ ] Performance improvement - [ ] Other ## How Has This Been Tested? - + + +**Test commands:** +```bash +# Example: bash scripts/run_ci_test.sh +``` + +**Test results:** +- [ ] Unit tests pass +- [ ] Integration tests pass (if applicable) +- [ ] Manual testing done (describe below) ## Checklist -- [ ] I have performed a self-review of my own code. -- [ ] I have formatted my own code using `./scripts/code_format.sh` before submitting. -- [ ] I have updated the documentation. -- [ ] I have added tests to prove my changes are effective. +- [ ] I have performed a self-review of my own code +- [ ] I have formatted my code using `./scripts/code_format.sh` +- [ ] I have run pre-commit on the files changed in this PR and all hooks pass +- [ ] I have updated the documentation (if applicable) +- [ ] I have added tests to prove my changes are effective +- [ ] For changes >500 LOC: I have filed an RFC issue + +## AI Assistance Disclosure + + + +- [ ] No AI tools were used +- [ ] AI tools were used (specify below) + + diff --git a/.github/workflows/_build-efa-wheel.yaml b/.github/workflows/_build-efa-wheel.yaml new file mode 100644 index 0000000000..730e7d87cc --- /dev/null +++ b/.github/workflows/_build-efa-wheel.yaml @@ -0,0 +1,223 @@ +name: _build-efa-wheel + +# Shared AWS EFA wheel build used by pull-request CI and release workflows. +# EFA hardware is not required to build: the distro libfabric is used for +# compilation, then auditwheel leaves libfabric/libefa to the system runtime. + +on: + workflow_call: + inputs: + variant: + # cuda | cuda13 | non-cuda + type: string + required: true + use-cuda: + type: boolean + required: true + python-versions: + # JSON array consumed by the called workflow's Python matrix. + type: string + required: true + build-profile: + # ci uses Ninja and CI diagnostics; release preserves the release make build. + type: string + required: true + cmake-args: + type: string + required: true + variant-flag: + # build_wheel.sh package variant, e.g. EFA_BUILD. + type: string + required: true + cuda-version: + type: string + default: '12.8.1' + torch-cuda-arch-list: + type: string + default: '' + artifact-prefix: + type: string + required: true + +env: + SCCACHE_GHA_ENABLED: "true" + +jobs: + build: + runs-on: ubuntu-22.04 + strategy: + matrix: + python-version: ${{ fromJSON(inputs.python-versions) }} + env: + BUILD_PROFILE: ${{ inputs.build-profile }} + CMAKE_ARGS: ${{ inputs.cmake-args }} + EFA_VARIANT: ${{ inputs.variant }} + USE_CUDA: ${{ inputs.use-cuda }} + VARIANT_FLAG: ${{ inputs.variant-flag }} + CUDA_VERSION: ${{ inputs.cuda-version }} + + steps: + - name: Validate build inputs + run: | + case "$BUILD_PROFILE" in + ci|release) ;; + *) echo "::error::Unknown EFA build profile: $BUILD_PROFILE"; exit 1 ;; + esac + + case "$EFA_VARIANT:$USE_CUDA:$VARIANT_FLAG" in + cuda:true:EFA_BUILD|cuda13:true:EFA_CU13_BUILD|non-cuda:false:EFA_NON_CUDA_BUILD) ;; + *) echo "::error::Inconsistent EFA variant inputs"; exit 1 ;; + esac + shell: bash + + - name: Configure CUDA architecture list + if: ${{ inputs.torch-cuda-arch-list != '' }} + env: + ARCH_LIST: ${{ inputs.torch-cuda-arch-list }} + run: echo "TORCH_CUDA_ARCH_LIST=$ARCH_LIST" >> "$GITHUB_ENV" + shell: bash + + - name: Checkout source + uses: actions/checkout@v4 + with: + persist-credentials: ${{ inputs.build-profile == 'release' }} + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Free up disk space + uses: ./.github/actions/free-disk-space + + - name: Install CUDA Toolkit + if: ${{ inputs.use-cuda }} + uses: Jimver/cuda-toolkit@v0.2.29 + with: + cuda: ${{ inputs.cuda-version }} + method: 'network' + sub-packages: '["nvcc", "nvrtc-dev"]' + non-cuda-sub-packages: '["libcusparse-dev", "libcublas-dev", "libcusolver-dev"]' + + - name: Run sccache-cache + uses: mozilla-actions/sccache-action@v0.0.9 + + - name: Configure sccache + uses: actions/github-script@v7 + with: + script: | + core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); + core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); + + - name: Install dependencies + run: | + sudo apt update -y + packages=(libfabric-dev libfabric1) + if [ "$BUILD_PROFILE" = ci ]; then + packages=(ninja-build "${packages[@]}") + fi + sudo apt install -y "${packages[@]}" + sudo bash -x dependencies.sh -y + if [ "$BUILD_PROFILE" = ci ]; then + df -h + fi + shell: bash + + - name: Configure project + run: | + generator=() + if [ "$BUILD_PROFILE" = ci ]; then + generator=(-G Ninja) + fi + + cuda_args=(-DUSE_CUDA=OFF) + if [ "$USE_CUDA" = true ]; then + cuda_args=( + -DUSE_CUDA=ON + -DCMAKE_EXE_LINKER_FLAGS=-L/usr/local/cuda/lib64/stubs + ) + fi + + read -r -a profile_args <<< "$CMAKE_ARGS" + mkdir build + cd build + cmake "${generator[@]}" .. \ + -DUSE_EFA=ON \ + -DLIBFABRIC_INCLUDE_DIR=/usr/include \ + -DLIBFABRIC_LIBRARY=/usr/lib/x86_64-linux-gnu/libfabric.so \ + -DENABLE_SCCACHE=ON \ + -DCMAKE_BUILD_TYPE=Release \ + "${profile_args[@]}" \ + "${cuda_args[@]}" + shell: bash + + - name: Build project + run: | + if [ "$USE_CUDA" = true ]; then + export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH + export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH + fi + + cd build + if [ "$BUILD_PROFILE" = release ]; then + make -j3 + sudo make install + else + cmake --build . + sudo cmake --install . + df -h + fi + shell: bash + + - name: Run sccache stat for check + if: ${{ env.SCCACHE_PATH != '' }} + shell: bash + run: ${SCCACHE_PATH} --show-stats + + - name: Generate Python version tag + id: python-tag + run: | + echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> "$GITHUB_OUTPUT" + shell: bash + + - name: Build Python wheel + run: | + export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib + export "$VARIANT_FLAG=1" + PYTHON_VERSION=${{ matrix.python-version }} \ + OUTPUT_DIR=dist-py${{ steps.python-tag.outputs.python_version_tag }} \ + ./scripts/build_wheel.sh + shell: bash + + - name: Verify libfabric is excluded from the wheel + run: | + WHL=$(ls mooncake-wheel/dist-py${{ steps.python-tag.outputs.python_version_tag }}/*.whl | head -1) + echo "Inspecting $WHL" + if unzip -l "$WHL" | grep -iE 'libfabric|libefa'; then + echo "::error::libfabric/libefa must NOT be bundled in the EFA wheel" + exit 1 + fi + echo "OK: libfabric/libefa correctly excluded (resolve to system EFA at runtime)" + shell: bash + + - name: Verify CUDA runtime dependency + if: ${{ inputs.use-cuda }} + run: | + WHL=$(ls mooncake-wheel/dist-py${{ steps.python-tag.outputs.python_version_tag }}/*.whl | head -1) + inspect_dir=$(mktemp -d) + unzip -q "$WHL" mooncake/engine.so -d "$inspect_dir" + cuda_major=${CUDA_VERSION%%.*} + if ! readelf -d "$inspect_dir/mooncake/engine.so" | \ + grep -Fq "Shared library: [libcudart.so.${cuda_major}]"; then + echo "::error::EFA wheel does not depend on libcudart.so.${cuda_major}" + readelf -d "$inspect_dir/mooncake/engine.so" | grep NEEDED + exit 1 + fi + echo "OK: EFA wheel depends on libcudart.so.${cuda_major}" + shell: bash + + - name: Upload Python wheel artifact + uses: actions/upload-artifact@v4 + with: + name: ${{ inputs.artifact-prefix }}-py${{ steps.python-tag.outputs.python_version_tag }} + path: mooncake-wheel/dist-py${{ steps.python-tag.outputs.python_version_tag }}/*.whl diff --git a/.github/workflows/_build-wheel.yaml b/.github/workflows/_build-wheel.yaml new file mode 100644 index 0000000000..21a0ed9e53 --- /dev/null +++ b/.github/workflows/_build-wheel.yaml @@ -0,0 +1,277 @@ +name: _build-wheel + +# Shared build for one supported wheel variant (matrixed over Python version). +# Callers choose the package variant and architecture; toolchain and CMake +# details live here so CI, nightly, pre-release, and release cannot drift. + +on: + workflow_call: + inputs: + variant: + # cuda | cuda13 | non-cuda + type: string + default: cuda + architecture: + # x86_64 | arm64 + type: string + default: x86_64 + python-versions: + type: string + default: '["3.10", "3.11", "3.12", "3.13"]' + artifact-prefix: + type: string + required: true + version-override: + # Optional: override the wheel version instead of deriving from git tag. + # Used by nightly builds to stamp date-based dev versions. + type: string + default: '' + +env: + SCCACHE_GHA_ENABLED: "true" + +jobs: + build: + runs-on: ${{ inputs.architecture == 'arm64' && 'ubuntu-22.04-arm' || 'ubuntu-22.04' }} + # manylinux2_28 keeps the x86 wheel's libstdc++/glibc floor compatible with + # RHEL/Rocky/Alma 8+. The CUDA image is also the toolchain for non-CUDA builds. + container: >- + ${{ inputs.architecture == 'arm64' && + (inputs.variant == 'cuda13' && 'pytorch/manylinuxaarch64-builder:cuda13.0' || 'pytorch/manylinuxaarch64-builder:cuda12.8') || + (inputs.variant == 'cuda13' && 'pytorch/manylinux2_28-builder:cuda13.0' || 'pytorch/manylinux2_28-builder:cuda12.8') }} + permissions: + contents: read + strategy: + matrix: + python-version: ${{ fromJSON(inputs.python-versions) }} + env: + BUILD_ARCHITECTURE: ${{ inputs.architecture }} + BUILD_VARIANT: ${{ inputs.variant }} + steps: + - name: Checkout source + uses: actions/checkout@v4 + + - name: Mark workspace safe for git + run: git config --global --add safe.directory '*' + + - name: Configure build profile + shell: bash + run: | + case "$BUILD_VARIANT:$BUILD_ARCHITECTURE" in + cuda:x86_64) + variant_flag='' + generator='' + arch_list='8.0;9.0' + ep_versions='2.11.0;2.12.0;2.12.1;2.13.0' + build_nvlink=true + cmake_args='-DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DUSE_INTRA_NVLINK=ON -DWITH_EP=ON -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release' + ;; + cuda13:x86_64) + variant_flag=CU13_BUILD + generator='' + arch_list='8.0;9.0;10.3' + ep_versions='2.11.0;2.12.0;2.12.1;2.13.0' + build_nvlink=true + cmake_args='-DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DUSE_INTRA_NVLINK=ON -DWITH_EP=ON -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release' + ;; + non-cuda:x86_64) + variant_flag=NON_CUDA_BUILD + generator='' + arch_list='' + ep_versions='' + build_nvlink=false + cmake_args='-DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=OFF -DWITH_EP=OFF -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release' + ;; + cuda:arm64) + variant_flag='' + generator=Ninja + arch_list='9.0' + ep_versions='2.11.0;2.12.0;2.12.1;2.13.0' + build_nvlink=false + cmake_args='-DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_CUDA=ON -DUSE_MNNVL=ON -DWITH_EP=ON -DWITH_STORE_RUST=OFF -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release' + ;; + cuda13:arm64) + variant_flag=CU13_BUILD + generator=Ninja + arch_list='9.0;10.3' + ep_versions='2.11.0;2.12.0;2.12.1;2.13.0' + build_nvlink=false + cmake_args='-DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_CUDA=ON -DUSE_MNNVL=ON -DWITH_EP=ON -DWITH_STORE_RUST=OFF -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release' + ;; + non-cuda:arm64) + variant_flag=NON_CUDA_BUILD + generator=Ninja + arch_list='' + ep_versions='' + build_nvlink=false + cmake_args='-DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=OFF -DWITH_EP=OFF -DSTORE_USE_ETCD=ON -DWITH_STORE_RUST=OFF -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release' + ;; + *) + echo "::error::Unsupported wheel profile: $BUILD_VARIANT/$BUILD_ARCHITECTURE" + exit 1 + ;; + esac + + { + echo "VARIANT_FLAG=$variant_flag" + echo "CMAKE_GEN=$generator" + echo "TORCH_CUDA_ARCH_LIST=$arch_list" + echo "EP_TORCH_VERSIONS_INPUT=$ep_versions" + echo "BUILD_NVLINK_ALLOCATOR=$build_nvlink" + echo "CMAKE_ARGS=$cmake_args" + } >> "$GITHUB_ENV" + + - name: Set version + run: | + if [ -n "${{ inputs.version-override }}" ]; then + echo "VERSION=${{ inputs.version-override }}" >> "$GITHUB_ENV" + else + echo "VERSION=${GITHUB_REF_NAME#v}" >> "$GITHUB_ENV" + fi + + - name: Patch wheel version in pyproject.toml + if: ${{ inputs.version-override != '' }} + run: | + sed -i "s/^version = .*/version = \"$VERSION\"/" mooncake-wheel/pyproject.toml + echo "Patched mooncake-wheel/pyproject.toml version to: $VERSION" + + - name: Select Python ${{ matrix.python-version }} from manylinux image + run: | + PYV_NODOT=$(echo "${{ matrix.python-version }}" | tr -d '.') + PYBIN="/opt/python/cp${PYV_NODOT}-cp${PYV_NODOT}/bin" + echo "$PYBIN" >> "$GITHUB_PATH" + "$PYBIN/pip" install --quiet "cmake<4" setuptools wheel + + - name: Run sccache-cache + uses: mozilla-actions/sccache-action@v0.0.9 + + - name: Configure sccache + uses: actions/github-script@v7 + with: + script: | + core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); + core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); + + - name: Configure project + run: | + SUDO=""; [ "$(id -u)" -ne 0 ] && SUDO="sudo" + if [ -n "$SUDO" ] && command -v apt-get >/dev/null 2>&1; then + $SUDO apt-get update -y || true + fi + $SUDO bash -x dependencies.sh -y + echo "/usr/local/go/bin" >> "$GITHUB_PATH" + gen=(); [ -n "$CMAKE_GEN" ] && gen=(-G "$CMAKE_GEN") + ep=(); [ -n "$EP_TORCH_VERSIONS_INPUT" ] && ep=(-DEP_TORCH_VERSIONS="$EP_TORCH_VERSIONS_INPUT") + mkdir -p build && cd build + # shellcheck disable=SC2086 + cmake "${gen[@]}" .. $CMAKE_ARGS "${ep[@]}" -DPython3_EXECUTABLE="$(which python3)" + + - name: Build project + env: + # Bound the nested Ninja build used by torch.utils.cpp_extension. + MAX_JOBS: "2" + run: | + for dir in /usr/local/cuda/lib64/stubs /usr/local/cuda/targets/*/lib/stubs; do + [ -d "$dir" ] && export LIBRARY_PATH="$dir:${LIBRARY_PATH:-}" + done + [ -d /usr/local/cuda ] && export CUDA_HOME=/usr/local/cuda + SUDO=""; [ "$(id -u)" -ne 0 ] && SUDO="sudo -E" + cd build + cmake --build . -j"$(nproc)" + $SUDO cmake --install . + + - name: Build nvlink_allocator.so + if: ${{ env.BUILD_NVLINK_ALLOCATOR == 'true' }} + run: | + export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH + if [ -d /usr/local/cuda/lib64/stubs ]; then + export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:${LD_LIBRARY_PATH:-} + export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:${LIBRARY_PATH:-} + fi + mkdir -p build/mooncake-transfer-engine/nvlink-allocator + cd mooncake-transfer-engine/nvlink-allocator + bash build.sh ../../build/mooncake-transfer-engine/nvlink-allocator/ + + - name: Run sccache stat for check + if: ${{ env.SCCACHE_PATH != '' }} + run: ${SCCACHE_PATH} --show-stats + + - name: Generate Python version tag + id: pytag + run: echo "tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> "$GITHUB_OUTPUT" + + - name: Build Python wheel + run: | + [ -d /usr/local/cuda ] && export CUDA_HOME=/usr/local/cuda + export LD_LIBRARY_PATH="${LD_LIBRARY_PATH:-}:/usr/local/lib" + variant=() + [ -n "$VARIANT_FLAG" ] && variant=("$VARIANT_FLAG=1") + env "${variant[@]}" \ + PYTHON_VERSION="${{ matrix.python-version }}" \ + OUTPUT_DIR="dist-py${{ steps.pytag.outputs.tag }}" \ + ./scripts/build_wheel.sh + env: + VERSION: ${{ env.VERSION }} + + - name: Smoke test repaired wheel + shell: bash + run: | + smoke_venv=$(mktemp -d) + python -m venv "$smoke_venv" + "$smoke_venv/bin/python" -m pip install --no-deps \ + mooncake-wheel/dist-py${{ steps.pytag.outputs.tag }}/*.whl + + if [ "${VARIANT_FLAG:-}" = "NON_CUDA_BUILD" ]; then + site_packages=$("$smoke_venv/bin/python" -c \ + 'import sysconfig; print(sysconfig.get_paths()["purelib"])') + cuda_dependency_found=false + for package_path in "$site_packages"/mooncake*; do + [ -e "$package_path" ] || continue + while IFS= read -r -d '' file; do + cuda_dependencies=$(readelf -d "$file" 2>/dev/null | grep -E \ + 'Shared library: \[(libcuda|libcudart|libcublas|libcufft|libcurand|libcusolver|libcusparse|libcufile|libcupti|libnvrtc|libnvJitLink|libnvToolsExt|libnvfatbin|libnvidia|libnccl)\.so' || true) + if [ -n "$cuda_dependencies" ]; then + echo "::error file=$file::Non-CUDA wheel artifact depends on CUDA" + echo "$cuda_dependencies" + cuda_dependency_found=true + fi + done < <(find "$package_path" -type f -print0) + done + if [ "$cuda_dependency_found" = true ]; then + exit 1 + fi + fi + + export LD_LIBRARY_PATH="/usr/local/lib:${LD_LIBRARY_PATH:-}" + for dir in /usr/local/cuda/lib64 \ + /usr/local/cuda/lib64/stubs \ + /usr/local/cuda/targets/*/lib \ + /usr/local/cuda/targets/*/lib/stubs; do + [ -d "$dir" ] && export LD_LIBRARY_PATH="$dir:$LD_LIBRARY_PATH" + done + + # CUDA builder images contain the link-time driver stub but not the + # real libcuda.so.1 supplied by an NVIDIA driver. Non-CUDA wheels must + # not receive this alias, so an accidental CUDA dependency still fails. + if [ "${VARIANT_FLAG:-}" != "NON_CUDA_BUILD" ]; then + cuda_stub= + for candidate in /usr/local/cuda/lib64/stubs/libcuda.so \ + /usr/local/cuda/targets/*/lib/stubs/libcuda.so; do + if [ -f "$candidate" ]; then + cuda_stub="$candidate" + break + fi + done + if [ -n "$cuda_stub" ]; then + ln -s "$cuda_stub" "$smoke_venv/libcuda.so.1" + export LD_LIBRARY_PATH="$smoke_venv:$LD_LIBRARY_PATH" + fi + fi + + "$smoke_venv/bin/mooncake_master" --version + + - name: Upload Python wheel artifact + uses: actions/upload-artifact@v4 + with: + name: ${{ inputs.artifact-prefix }}-py${{ steps.pytag.outputs.tag }} + path: mooncake-wheel/dist-py${{ steps.pytag.outputs.tag }}/*.whl diff --git a/.github/workflows/_publish-wheel.yaml b/.github/workflows/_publish-wheel.yaml new file mode 100644 index 0000000000..c5bab698dd --- /dev/null +++ b/.github/workflows/_publish-wheel.yaml @@ -0,0 +1,49 @@ +name: _publish-wheel + +# Shared publish tail: collect this run's wheels, attach to the GitHub Release, +# and upload to PyPI. Used by the Release workflows. + +on: + workflow_call: + inputs: + artifact-pattern: + type: string + required: true + secrets: + pypi-token: + required: false + +jobs: + publish: + runs-on: ubuntu-22.04 + permissions: + contents: write + id-token: write + steps: + - name: Checkout source + uses: actions/checkout@v4 + + - name: Download all wheel artifacts + uses: actions/download-artifact@v4 + with: + path: mooncake-wheel/dist-all + pattern: ${{ inputs.artifact-pattern }} + + - name: Prepare wheels for release + run: | + mkdir -p mooncake-wheel/dist-release + find mooncake-wheel/dist-all -name "*.whl" -exec cp {} mooncake-wheel/dist-release/ \; + echo "Collected wheels for release:" + ls -la mooncake-wheel/dist-release/ + + - name: Upload wheels to GitHub Release + uses: softprops/action-gh-release@v1 + with: + files: mooncake-wheel/dist-release/*.whl + + - name: Publish package to PyPI + if: ${{ github.repository == 'kvcache-ai/Mooncake' }} + uses: pypa/gh-action-pypi-publish@release/v1 + with: + packages-dir: mooncake-wheel/dist-release/ + password: ${{ secrets.pypi-token }} diff --git a/.github/workflows/assistant.yml b/.github/workflows/assistant.yml new file mode 100644 index 0000000000..b04eeae3ba --- /dev/null +++ b/.github/workflows/assistant.yml @@ -0,0 +1,89 @@ +name: Qoder Assistant + +on: + issue_comment: + types: [created] + pull_request_review_comment: + types: [created] + +jobs: + qoder-assistant: + if: | + contains(github.event.comment.body, '@qoder') && + !endsWith(github.event.comment.user.login, '[bot]') + runs-on: ubuntu-latest + permissions: + contents: read + issues: write + pull-requests: write + id-token: write + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Build Arguments + id: build_args + run: | + ARGS="REPO: ${{ github.repository }} + REQUEST_SOURCE: ${{ github.event_name }} + THREAD_ID: ${{ github.event.comment.node_id }} + COMMENT_ID: ${{ github.event.comment.id }} + AUTHOR: ${{ github.event.comment.user.login }} + BODY: ${{ github.event.comment.body }} + URL: ${{ github.event.comment.html_url }} + IS_PR: ${{ github.event.issue.pull_request != null || github.event_name == 'pull_request_review_comment' }} + ISSUE_OR_PR_NUMBER: ${{ github.event.issue.number || github.event.pull_request.number }}" + + if [ -n "${{ github.event.comment.pull_request_review_id }}" ]; then + ARGS="$ARGS + REVIEW_ID: ${{ github.event.comment.pull_request_review_id }}" + fi + + if [ -n "${{ github.event.comment.in_reply_to_id }}" ]; then + ARGS="$ARGS + REPLY_TO_COMMENT_ID: ${{ github.event.comment.in_reply_to_id }}" + fi + + echo "args<> $GITHUB_OUTPUT + echo "$ARGS" >> $GITHUB_OUTPUT + echo "EOF" >> $GITHUB_OUTPUT + + - name: Run Qoder Assistant + uses: QoderAI/qoder-action@v0 + with: + qoder_personal_access_token: ${{ secrets.QODER_PERSONAL_ACCESS_TOKEN }} + prompt: | + /assistant + ${{ steps.build_args.outputs.args }} + NO_AUTO_APPROVE: true + Never approve or request changes on pull requests. If submitting a review via mcp__qoder_github__submit_pending_pull_request_review, use event type "COMMENT" only. DO NOT use "APPROVE" or "REQUEST_CHANGES". + + - name: Dismiss accidental Qoder approval + if: always() && (github.event.issue.pull_request != null || github.event_name == 'pull_request_review_comment') + env: + GH_TOKEN: ${{ github.token }} + run: | + set -euo pipefail + PR_NUMBER="${{ github.event.issue.number || github.event.pull_request.number }}" + REPO="${{ github.repository }}" + + mapfile -t REVIEW_IDS < <( + gh api "repos/${REPO}/pulls/${PR_NUMBER}/reviews" \ + --jq '.[] | select(.user.login == "qoderai[bot]" and .state == "APPROVED") | .id' + ) + + if [ "${#REVIEW_IDS[@]}" -eq 0 ]; then + echo "No Qoder APPROVED reviews to dismiss." + exit 0 + fi + + for review_id in "${REVIEW_IDS[@]}"; do + echo "Dismissing Qoder approval review ${review_id}..." + gh api \ + --method PUT \ + "repos/${REPO}/pulls/${PR_NUMBER}/reviews/${review_id}/dismissals" \ + -f message="Automated Qoder approvals are disabled; a human maintainer must approve this pull request." + done diff --git a/.github/workflows/cancel-ci.yml b/.github/workflows/cancel-ci.yml new file mode 100644 index 0000000000..7d322facd0 --- /dev/null +++ b/.github/workflows/cancel-ci.yml @@ -0,0 +1,169 @@ +name: Cancel Queued and Running CI + +# Manually cancel every unfinished (queued / running / waiting / pending) +# workflow run in the repo. Useful for draining runners after a bad push, +# a runaway matrix, or before urgent maintenance. +# +# Unlike the per-workflow `concurrency: cancel-in-progress` settings (which +# only cancel a run superseded by a newer run on the same ref), this cancels +# runs across all workflows and refs in one shot. +# +# Adapted from sglang's cancel-pr-workflows-on-close.yml cancellation logic: +# fetch recent runs, keep the non-terminal ones, force-cancel runs stuck +# behind approval gates, and run a second pass to catch runs still +# materializing during the first. +# +# Listing note: query params (branch, per_page, page) are placed directly in +# the API path rather than via `gh api -f key=value`. `gh api` treats -f/-F +# fields as a request *body* and, unless the method is pinned, flips GET to +# POST -- either way the value never reaches the query string, so a status/ +# branch filter silently matches nothing. Status is filtered client-side in +# jq to avoid that class of bug entirely. + +on: + workflow_dispatch: + inputs: + branch: + description: 'Only cancel runs on this branch (leave empty to cancel across all branches)' + required: false + type: string + default: '' + dry_run: + description: 'List what would be cancelled without cancelling' + required: false + type: boolean + default: false + +permissions: + actions: write # Needed to cancel runs + contents: read # Needed to read repo info + +# Do not let two drain operations fight each other. +concurrency: + group: cancel-ci + cancel-in-progress: false + +jobs: + cancel: + runs-on: ubuntu-latest + steps: + - name: Cancel unfinished runs + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + REPO: ${{ github.repository }} + BRANCH: ${{ github.event.inputs.branch }} + DRY_RUN: ${{ github.event.inputs.dry_run }} + shell: bash + run: | + set -euo pipefail + + # Non-terminal run statuses. GitHub reports in-flight runs as one of + # these; everything else (completed) is terminal. + UNFINISHED_RE='^(queued|in_progress|waiting|pending|requested|action_required)$' + + echo "==================================================" + echo " Cancel Queued and Running CI" + echo "==================================================" + echo "Repo: $REPO" + if [ -n "$BRANCH" ]; then + echo "Scope: branch '$BRANCH' only" + else + echo "Scope: ALL branches" + fi + if [ "$DRY_RUN" = "true" ]; then + echo "" + echo " ⚠️ DRY RUN — runs will be LISTED but NOT cancelled." + echo " ⚠️ Re-run with dry_run unchecked to actually cancel." + fi + echo "==================================================" + echo "" + + FAILURES=0 + + # List unfinished runs. Params go in the path (see header note); status + # is filtered client-side. Only the most-recent pages can hold in-flight + # runs, so a few pages by created-desc order is plenty and stays bounded. + list_unfinished() { + local ids="" page batch + for page in 1 2 3 4 5; do + local path="repos/$REPO/actions/runs?per_page=100&page=$page" + if [ -n "$BRANCH" ]; then + path="$path&branch=$BRANCH" + fi + batch="" + for attempt in 1 2 3; do + if batch=$(gh api -X GET "$path" \ + --jq ".workflow_runs[] + | select(.id != $GITHUB_RUN_ID) + | select(.status | test(\"$UNFINISHED_RE\")) + | \"\(.id)\t\(.status)\t\(.name)\""); then + break + fi + batch="" + # Runs in a $(...) subshell, so an error annotation is the only + # signal that survives; the retries make a hard failure rare. + if [ "$attempt" -eq 3 ]; then + echo "::error::Listing runs (page $page) failed after 3 attempts" >&2 + else + sleep 5 + fi + done + ids="$ids"$'\n'"$batch" + done + echo "$ids" | sed '/^$/d' | sort -u + } + + # Pass 2 catches runs still materializing during pass 1. + for pass in 1 2; do + rows=$(list_unfinished) + + if [ -z "$rows" ]; then + echo "Pass $pass: no unfinished runs found" + break + fi + echo "Pass $pass: found $(echo "$rows" | wc -l | tr -d ' ') unfinished run(s):" + echo "$rows" | while IFS=$'\t' read -r rid rstatus rname; do + echo " • $rid [$rstatus] $rname" + done + echo "" + + run_ids=$(echo "$rows" | cut -f1) + for run_id in $run_ids; do + run_url="https://github.com/$REPO/actions/runs/$run_id" + if [ "$DRY_RUN" = "true" ]; then + echo " [dry-run] would cancel $run_url" + continue + fi + echo "Cancelling $run_url" + if gh run cancel "$run_id" --repo "$REPO" 2>/dev/null; then + continue + fi + # Plain cancel fails for runs stuck behind approval / + # deployment protection rules; force-cancel handles those. + if gh api -X POST "repos/$REPO/actions/runs/$run_id/force-cancel" >/dev/null 2>&1; then + echo " force-cancelled" + continue + fi + # Finished between listing and cancelling is fine. + state=$(gh api "repos/$REPO/actions/runs/$run_id" --jq '.status' 2>/dev/null || echo "unknown") + if [ "$state" = "completed" ]; then + echo " already finished, nothing to cancel" + else + echo "::error::Failed to cancel run $run_id (status: $state)" + FAILURES=1 + fi + done + + if [ "$DRY_RUN" = "true" ]; then + echo "Pass $pass: dry run, skipping second pass" + break + fi + + if [ "$pass" -eq 1 ]; then + sleep 20 + fi + done + + echo "" + echo "✅ Done" + exit "$FAILURES" diff --git a/.github/workflows/ci-on-label.yml b/.github/workflows/ci-on-label.yml new file mode 100644 index 0000000000..f4d5929ca6 --- /dev/null +++ b/.github/workflows/ci-on-label.yml @@ -0,0 +1,53 @@ +name: Retrigger CI on run-ci label + +# Same-SHA retrigger for Build & Test. This is a separate workflow so labels +# other than `run-ci` (especially `run-e2e-ci`) cannot start or cancel +# `.github/workflows/ci.yml`. +# +# pull_request_target is required so fork PRs can rerun Actions. This +# workflow only calls the GitHub API; it does not check out PR code. +on: + pull_request_target: + branches: + - "main" + - "release/**" + types: [labeled] + +permissions: + actions: write + contents: read + +jobs: + retrigger: + if: > + github.event.label.name == 'run-ci' && + github.actor != 'github-actions[bot]' + runs-on: ubuntu-latest + steps: + - name: Re-run Build & Test for this SHA + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + REPO: ${{ github.repository }} + SHA: ${{ github.event.pull_request.head.sha }} + shell: bash + run: | + set -euo pipefail + + run_json=$(gh api \ + "repos/${REPO}/actions/workflows/ci.yml/runs?head_sha=${SHA}&per_page=20") + run_id=$(echo "$run_json" | jq -r '.workflow_runs[0].id // empty') + status=$(echo "$run_json" | jq -r '.workflow_runs[0].status // empty') + + if [ -z "$run_id" ]; then + echo "No Build & Test run found for SHA ${SHA}." + echo "Open or push to the PR first so ci.yml has a run to rerun." + exit 1 + fi + + echo "Matched workflow run ${run_id} (status=${status})" + if [ "$status" != "completed" ]; then + echo "Build & Test is still ${status}; not starting a duplicate." + exit 0 + fi + + gh run rerun "$run_id" --repo "$REPO" diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3f68759e64..a6cdb72374 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -2,49 +2,96 @@ name: 'Build & Test (Linux)' on: push: - branches: [ "main" ] + branches: + - "main" + - "release/**" pull_request: - branches: [ "main" ] - types: [opened, synchronize, reopened, labeled] - + branches: + - "main" + - "release/**" + # `labeled` is intentionally omitted. Auto-labeler already applies + # `run-ci`, so any new label (including `run-e2e-ci`) would retrigger + # this whole workflow and cancel the in-progress run. Same-SHA + # retrigger via the `run-ci` label lives in ci-on-label.yml. + types: [opened, synchronize, reopened] + 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: + # Build the artifact tested below through the exact release wheel path. The + # Ubuntu jobs remain consumer tests, not an alternate wheel build environment. + build-wheel: + needs: [spell-check, clang-format, check-paths] + if: &run-ci-for-source-changes >- + (needs.check-paths.outputs.should-run-downstream == 'true' || + github.event_name == 'workflow_dispatch') && + (github.event_name == 'push' || + github.event_name == 'workflow_dispatch' || + (github.event_name == 'pull_request' && + (github.event.action != 'labeled' || + contains(github.event.pull_request.labels.*.name, 'run-ci')))) + uses: ./.github/workflows/_build-wheel.yaml + with: + python-versions: '["3.10", "3.12"]' + artifact-prefix: mooncake-wheel-ci + version-override: 0.0.0.dev0 + + test-wheel-ubuntu: + needs: [build-wheel] if: >- - github.event_name == 'push' || - github.event.action == 'opened' || - contains(github.event.pull_request.labels.*.name, 'run-ci') - runs-on: ubuntu-22.04 + needs.build-wheel.result == 'success' && + (github.event_name == 'push' || + github.event_name == 'workflow_dispatch' || + (github.event_name == 'pull_request' && + (github.event.action != 'labeled' || + contains(github.event.pull_request.labels.*.name, 'run-ci')))) strategy: matrix: - python-version: ['3.10', '3.12'] - env: - CI: "true" - SCCACHE_GHA_ENABLED: "true" - + include: + - ubuntu-version: ubuntu-22.04 + python-version: '3.12' + - ubuntu-version: ubuntu-24.04 + python-version: '3.10' + 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 with: python-version: ${{ matrix.python-version }} - - name: Install and start etcd + - name: Generate Python version tag + id: generate_tag_test run: | - wget https://github.com/etcd-io/etcd/releases/download/v3.6.1/etcd-v3.6.1-linux-amd64.tar.gz - tar xzf etcd-v3.6.1-linux-amd64.tar.gz - sudo mv etcd-v3.6.1-linux-amd64/etcd* /usr/local/bin/ - etcd --advertise-client-urls http://127.0.0.1:2379 --listen-client-urls http://127.0.0.1:2379 & - sleep 3 # Give etcd time to start - etcdctl --endpoints=http://127.0.0.1:2379 endpoint health + echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT shell: bash - - name: Free up disk space + - name: Download wheel artifact + uses: actions/download-artifact@v4 + with: + name: mooncake-wheel-ci-py${{ steps.generate_tag_test.outputs.python_version_tag }} + path: mooncake-wheel/dist + + - name: Verify wheel file exists run: | - sudo rm -rf /usr/share/dotnet - sudo rm -rf /opt/ghc - sudo rm -rf /opt/hostedtoolcache/CodeQL + ls -la mooncake-wheel/dist/ + if [ ! -f mooncake-wheel/dist/*.whl ]; then + echo "ERROR: No wheel file found in mooncake-wheel/dist/" + exit 1 + fi + shell: bash + + - name: Free up disk space + uses: ./.github/actions/free-disk-space - name: Install CUDA Toolkit uses: Jimver/cuda-toolkit@v0.2.24 @@ -52,220 +99,252 @@ jobs: cuda: '12.8.1' linux-local-args: '["--toolkit"]' method: 'network' - sub-packages: '["nvcc"]' - - - name: Install coverage tools - run: | - sudo apt-get update - sudo apt-get install -y lcov gcovr - - - name: Set up coverage compilation flags - run: | - echo "Setting up coverage compilation flags..." - echo "CXXFLAGS=--coverage" >> $GITHUB_ENV - echo "CFLAGS=--coverage" >> $GITHUB_ENV - echo "LDFLAGS=--coverage" >> $GITHUB_ENV - shell: bash - - - name: Run sccache-cache - uses: mozilla-actions/sccache-action@v0.0.9 - - - name: Configure sccache - uses: actions/github-script@v7 - with: - script: | - core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); - core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); - - - name: Run sccache stat for check - shell: bash - run: ${SCCACHE_PATH} --show-stats - - name: Configure project with coverage support - run: | - sudo apt update -y - sudo bash -x dependencies.sh -y - mkdir build - cd build - 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 + - name: Run installation test script run: | - cd build - make -j4 - sudo make install + bash scripts/test_installation.sh shell: bash - - name: Build nvlink_allocator.so + - name: Run all PG CPU tests with latest PyTorch + if: matrix.ubuntu-version == 'ubuntu-22.04' && matrix.python-version == '3.12' + env: + MC_FORCE_TCP: "true" run: | - mkdir -p build/mooncake-transfer-engine/nvlink-allocator - cd mooncake-transfer-engine/nvlink-allocator - export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH - bash build.sh ../../build/mooncake-transfer-engine/nvlink-allocator/ + source test_env/bin/activate + export LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:/usr/local/lib + python -m unittest discover -s mooncake-pg/tests -k CPU -v shell: bash - - name: Start Metadata Server + - name: Run Store API and RPC smoke tests + if: matrix.ubuntu-version == 'ubuntu-22.04' && matrix.python-version == '3.12' + 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: | - cd mooncake-transfer-engine/example/http-metadata-server-python - pip install aiohttp - python ./bootstrap_server.py & - shell: bash + source test_env/bin/activate + export LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:/usr/local/lib + python -m pip install safetensors - - 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 make test -j ARGS="-V" - shell: bash + reserved_ports=$(sysctl -n net.ipv4.ip_local_reserved_ports) + sudo sysctl -w \ + "net.ipv4.ip_local_reserved_ports=${reserved_ports:+$reserved_ports,}50052" - - name: Generate coverage report - id: coverage - run: | - cd build - echo "=== Starting coverage report generation ===" - echo "Current directory: $(pwd)" - - echo "=== Looking for .gcda files ===" - find . -name "*.gcda" 2>/dev/null | head -10 || echo "No .gcda files found" - - echo "=== Running lcov ===" - lcov --capture --directory . --output-file coverage.info 2>&1 || { - echo "WARNING: lcov failed to capture coverage data" - echo "Creating minimal lcov-compliant coverage file to allow CI to continue" - echo "TN:dummy" > coverage.filtered.info - echo "SF:/dev/null" >> coverage.filtered.info - echo "DA:0,0" >> coverage.filtered.info - echo "end_of_record" >> coverage.filtered.info - echo "coverage_failed=true" >> $GITHUB_OUTPUT - exit 0 # Exit successfully, do not block CI + mkdir -p /tmp/mooncake_storage + mooncake_master \ + --default_kv_lease_ttl=500 \ + --eviction_high_watermark_ratio=0.95 \ + --cluster_id=ci_store_api_smoke \ + --port 50051 \ + --enable_http_metadata_server=true \ + >"$RUNNER_TEMP/mooncake-master.log" 2>&1 & + master_pid=$! + rpc_server_pid="" + + cleanup() { + if [ -n "$rpc_server_pid" ]; then + kill "$rpc_server_pid" 2>/dev/null || true + wait "$rpc_server_pid" 2>/dev/null || true + fi + kill "$master_pid" 2>/dev/null || true + wait "$master_pid" 2>/dev/null || true } + trap cleanup EXIT + + master_ready=false + for _ in {1..50}; do + if ! kill -0 "$master_pid" 2>/dev/null; then + cat "$RUNNER_TEMP/mooncake-master.log" + echo "::error::mooncake_master exited before becoming ready" + exit 1 + fi + if ss -H -ltn 'sport = :50051' | grep -q . && \ + ss -H -ltn 'sport = :8080' | grep -q .; then + master_ready=true + break + fi + sleep 0.1 + done + if [ "$master_ready" != true ]; then + cat "$RUNNER_TEMP/mooncake-master.log" + echo "::error::mooncake_master did not become ready within 5 seconds" + exit 1 + fi - echo "=== Processing coverage data ===" - lcov --remove coverage.info '/usr/*' '*/test/*' '*/third_party/*' --output-file coverage.filtered.info 2>&1 || true - - echo "=== Generating HTML report ===" - genhtml coverage.filtered.info --output-directory coverage_report 2>&1 || echo "genhtml failed, continuing..." - - echo "=== Coverage summary ===" - lcov --list coverage.filtered.info 2>&1 || echo "lcov list failed" - - echo "=== Coverage report generation completed ===" - shell: bash - - - name: Upload coverage to Codecov - uses: codecov/codecov-action@v4 - with: - files: build/coverage.filtered.info - flags: unittests - name: code-coverage-report - token: ${{ secrets.CODECOV_TOKEN }} - fail_ci_if_error: false - continue-on-error: true - - - name: Check coverage status - if: always() - run: | - if [ "${{ steps.coverage.outputs.coverage_failed }}" = "true" ]; then - echo "⚠️ Coverage collection failed but CI continued" - echo "::warning::Code coverage collection failed. Please check the build logs." - else - echo "✅ Coverage collected successfully" + python scripts/test_tensor_api.py -n 1 + python scripts/test_async_store.py + python scripts/test_copy_move_api.py + python -m unittest mooncake-wheel.tests.test_safetensor_functions + python scripts/test_drain_http_api.py --timeout-sec 90 + + python -u mooncake-transfer-engine/tests/rpc_communicator_test.py \ + server --url 127.0.0.1:9004 --data-size 1 \ + >"$RUNNER_TEMP/rpc-server.log" 2>&1 & + rpc_server_pid=$! + rpc_ready=false + for _ in {1..50}; do + if ! kill -0 "$rpc_server_pid" 2>/dev/null; then + cat "$RUNNER_TEMP/rpc-server.log" + echo "::error::RPC communicator server exited before becoming ready" + exit 1 + fi + if ss -H -ltn 'sport = :9004' | grep -q .; then + rpc_ready=true + break + fi + sleep 0.1 + done + if [ "$rpc_ready" != true ]; then + cat "$RUNNER_TEMP/rpc-server.log" + echo "::error::RPC communicator server did not become ready within 5 seconds" + exit 1 fi - - name: Generate Python version tag - id: generate_tag_build - run: | - echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT + client_rc=0 + timeout 10 python -u \ + mooncake-transfer-engine/tests/rpc_communicator_test.py \ + client --url 127.0.0.1:9004 --threads 2 --data-size 1 \ + >"$RUNNER_TEMP/rpc-client.log" 2>&1 || \ + client_rc=$? + cat "$RUNNER_TEMP/rpc-client.log" + if [ "$client_rc" -ne 0 ] && [ "$client_rc" -ne 124 ]; then + echo "::error::RPC communicator client failed with exit code $client_rc" + exit "$client_rc" + fi + if ! grep -q '^bandwidth:' "$RUNNER_TEMP/rpc-client.log"; then + cat "$RUNNER_TEMP/rpc-server.log" + echo "::error::RPC communicator did not complete a successful transfer" + exit 1 + fi shell: bash - # In CI, build_wheel.sh removes build/ to free disk (CI=true); set FREE_BUILD_DIR=1 locally to enable. - - name: Build Python wheel + - name: Run SSD offload and promotion end-to-end tests + if: matrix.ubuntu-version == 'ubuntu-22.04' && matrix.python-version == '3.12' run: | - PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_build.outputs.python_version_tag }} ./scripts/build_wheel.sh - shell: bash - - - name: Upload wheel for ZMQ test job - uses: actions/upload-artifact@v4 - with: - name: wheel-build-py${{ steps.generate_tag_build.outputs.python_version_tag }} - path: mooncake-wheel/dist-py${{ steps.generate_tag_build.outputs.python_version_tag }}/*.whl - - build-musa: - if: >- - 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 - - - name: Mark repository as safe - run: git config --global --add safe.directory $GITHUB_WORKSPACE - shell: bash + source test_env/bin/activate + export LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:/usr/local/lib + + metadata_pid="" + master_pid="" + cleanup() { + status=$? + trap - EXIT + if [ -n "$master_pid" ]; then + kill "$master_pid" 2>/dev/null || true + wait "$master_pid" 2>/dev/null || true + fi + if [ -n "$metadata_pid" ]; then + kill "$metadata_pid" 2>/dev/null || true + wait "$metadata_pid" 2>/dev/null || true + fi + if [ "$status" -ne 0 ]; then + cat "$RUNNER_TEMP/mooncake-ssd-master.log" 2>/dev/null || true + cat "$RUNNER_TEMP/mooncake-promotion-master.log" 2>/dev/null || true + cat "$RUNNER_TEMP/mooncake-metadata.log" 2>/dev/null || true + fi + rm -rf /tmp/mooncake_ci_ssd_offload \ + /tmp/mooncake_ci_promotion + exit "$status" + } + trap cleanup EXIT + + wait_for_port() { + local pid=$1 + local port=$2 + local label=$3 + for _ in {1..50}; do + if ! kill -0 "$pid" 2>/dev/null; then + echo "::error::$label exited before listening on port $port" + return 1 + fi + if ss -H -ltn "sport = :$port" | grep -q .; then + return 0 + fi + sleep 0.1 + done + echo "::error::$label did not listen on port $port within 5 seconds" + return 1 + } - - name: Configure project - run: | - apt update -y - bash -x dependencies.sh -y - mkdir build - cd build - 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 + stop_master() { + kill "$master_pid" 2>/dev/null || true + wait "$master_pid" 2>/dev/null || true + master_pid="" + } - - name: Build project - run: | - cd build - source ~/.bashrc - make -j4 - make install - shell: bash + mooncake_http_metadata_server --port 8080 \ + >"$RUNNER_TEMP/mooncake-metadata.log" 2>&1 & + metadata_pid=$! + wait_for_port "$metadata_pid" 8080 "metadata server" - test-wheel-ubuntu: - needs: build-flags - strategy: - matrix: - ubuntu-version: [ubuntu-22.04, ubuntu-24.04] - python-version: ['3.10', '3.12'] - runs-on: ${{ matrix.ubuntu-version }} + mkdir -p /tmp/mooncake_ci_ssd_offload + mooncake_master \ + --default_kv_lease_ttl=500 \ + --root_fs_dir=/tmp/mooncake_ci_ssd_offload \ + >"$RUNNER_TEMP/mooncake-ssd-master.log" 2>&1 & + master_pid=$! + wait_for_port "$master_pid" 50051 "SSD offload master" + MC_METADATA_SERVER=http://127.0.0.1:8080/metadata \ + DEFAULT_KV_LEASE_TTL=500 \ + python mooncake-wheel/tests/test_ssd_offload_in_evict.py + stop_master + + mkdir -p /tmp/mooncake_ci_promotion + mooncake_master \ + --default_kv_lease_ttl=500 \ + --root_fs_dir=/tmp/mooncake_ci_promotion \ + --enable_offload=true \ + --offload_on_evict=true \ + --promotion_on_hit=true \ + --promotion_admission_threshold=1 \ + --promotion_max_per_heartbeat=16 \ + >"$RUNNER_TEMP/mooncake-promotion-master.log" 2>&1 & + master_pid=$! + wait_for_port "$master_pid" 50051 "promotion master" + MC_METADATA_SERVER=http://127.0.0.1:8080/metadata \ + DEFAULT_KV_LEASE_TTL=500 \ + MOONCAKE_OFFLOAD_FILE_STORAGE_PATH=/tmp/mooncake_ci_promotion \ + MOONCAKE_OFFLOAD_HEARTBEAT_INTERVAL_SECONDS=2 \ + MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT=10 \ + MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES=10485760 \ + python mooncake-wheel/tests/test_promotion_on_hit.py + stop_master + shell: bash + + unit-tests: + name: CTest unit tests (Python 3.12) + needs: [check-paths] + if: *run-ci-for-source-changes + runs-on: ubuntu-22.04 + env: + CI: "true" + SCCACHE_GHA_ENABLED: "true" steps: - uses: actions/checkout@v4 + with: + persist-credentials: false - - name: Set up Python ${{ matrix.python-version }} + - name: Set up Python 3.12 uses: actions/setup-python@v5 with: - python-version: ${{ matrix.python-version }} + python-version: '3.12' - - name: Generate Python version tag - id: generate_tag_test - run: | - echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT - shell: bash + - name: Install Rust toolchain + uses: dtolnay/rust-toolchain@stable - - name: Download wheel artifact - uses: actions/download-artifact@v4 - with: - name: mooncake-wheel-ubuntu-py${{ steps.generate_tag_test.outputs.python_version_tag }} - path: mooncake-wheel/dist - - - name: Verify wheel file exists + - name: Install and start etcd run: | - ls -la mooncake-wheel/dist/ - if [ ! -f mooncake-wheel/dist/*.whl ]; then - echo "ERROR: No wheel file found in mooncake-wheel/dist/" - exit 1 - fi + wget -q https://github.com/etcd-io/etcd/releases/download/v3.6.1/etcd-v3.6.1-linux-amd64.tar.gz + tar xzf etcd-v3.6.1-linux-amd64.tar.gz + sudo mv etcd-v3.6.1-linux-amd64/etcd* /usr/local/bin/ + etcd --advertise-client-urls http://127.0.0.1:2379 --listen-client-urls http://127.0.0.1:2379 & + sleep 3 + ETCDCTL_API=3 etcdctl --endpoints=http://127.0.0.1:2379 endpoint health shell: bash - name: Free up disk space - run: | - sudo rm -rf /usr/share/dotnet - sudo rm -rf /opt/ghc - sudo rm -rf /opt/hostedtoolcache/CodeQL - sudo rm -rf /usr/local/lib/android - df -h + uses: ./.github/actions/free-disk-space - name: Install CUDA Toolkit uses: Jimver/cuda-toolkit@v0.2.24 @@ -273,117 +352,128 @@ jobs: cuda: '12.8.1' linux-local-args: '["--toolkit"]' method: 'network' + sub-packages: '["nvcc"]' - - name: Run installation test script + - name: Install build utilities and dependencies run: | - bash scripts/test_installation.sh + sudo apt-get update + sudo apt-get install -y ninja-build + sudo bash -x dependencies.sh -y + echo "/usr/local/go/bin" >> "$GITHUB_PATH" shell: bash - - name: Start metadata server - run: | - source test_env/bin/activate - mooncake_http_metadata_server --port 8080 & - shell: bash + - name: Run sccache-cache + uses: mozilla-actions/sccache-action@v0.0.9 - - name: Run tests with ssd + - name: Configure sccache + uses: actions/github-script@v7 + with: + script: | + core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); + core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); + + - name: Configure project with unit tests run: | - 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 - deactivate + cmake -S . -B build -G Ninja \ + -DUSE_HTTP=ON \ + -DUSE_CXL=ON \ + -DUSE_UB=ON \ + -DUSE_ETCD=ON \ + -DUSE_CUDA=ON \ + -DWITH_P2P_STORE=ON \ + -DSTORE_USE_ETCD=ON \ + -DCMAKE_BUILD_TYPE=Release \ + -DBUILD_UNIT_TESTS=ON \ + -DENABLE_SCCACHE=ON shell: bash - - name: Start Mooncake Master + - name: Configure CUDA driver runtime run: | - source test_env/bin/activate - mkdir -p /tmp/mooncake_storage - mooncake_master \ - --eviction_high_watermark_ratio=0.95 \ - --cluster_id=ci_test_cluster \ - --port 50051 & - sleep 3 + cuda_driver_library=$(sed -n \ + 's/^CUDA_cuda_driver_LIBRARY:FILEPATH=//p' build/CMakeCache.txt) + if [ -z "$cuda_driver_library" ] || [ ! -f "$cuda_driver_library" ]; then + echo "::error::CMake did not resolve the CUDA driver library" + exit 1 + fi + + cuda_driver_dir=$(dirname "$cuda_driver_library") + if [ ! -e "$cuda_driver_dir/libcuda.so.1" ]; then + sudo ln -s "$(basename "$cuda_driver_library")" \ + "$cuda_driver_dir/libcuda.so.1" + fi + echo "LIBRARY_PATH=$cuda_driver_dir:${LIBRARY_PATH:-}" >> "$GITHUB_ENV" + echo "LD_LIBRARY_PATH=$cuda_driver_dir:${LD_LIBRARY_PATH:-}" >> "$GITHUB_ENV" shell: bash - - name: Run Python Tensor API Performance 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" + - name: Build project and unit tests run: | - source test_env/bin/activate - python scripts/test_tensor_api.py -n 1 + cmake --build build --parallel 128 + sudo -E cmake --install build shell: bash - - name: Run Python Async API Test (CI check) + - name: Run Mooncake Store Rust unit tests 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" + MOONCAKE_BUILD_DIR: ${{ github.workspace }}/build + MOONCAKE_STORE_LIB_DIR: ${{ github.workspace }}/build/mooncake-store/src + MOONCAKE_STORE_INCLUDE_DIR: ${{ github.workspace }}/mooncake-store/include run: | - source test_env/bin/activate - python scripts/test_async_store.py + export LD_LIBRARY_PATH="$GITHUB_WORKSPACE/build/mooncake-asio:$GITHUB_WORKSPACE/build/mooncake-store/src:$GITHUB_WORKSPACE/build/mooncake-store/src/cachelib_memory_allocator:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base:$GITHUB_WORKSPACE/build/mooncake-common/etcd:/usr/local/lib:${LD_LIBRARY_PATH:-}" + cargo test --manifest-path mooncake-store/rust/Cargo.toml --lib shell: bash - - name: Test Mooncake Copy/Move API - 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" + - name: Start Metadata Server run: | - source test_env/bin/activate - python scripts/test_copy_move_api.py + cd mooncake-transfer-engine/example/http-metadata-server-python + pip install aiohttp + python ./bootstrap_server.py & + sleep 2 shell: bash - - name: Run RPC Communicator Bandwidth Test - run: | - source test_env/bin/activate - python mooncake-transfer-engine/tests/rpc_communicator_test.py server --url 127.0.0.1:9004 --data-size 1 & - SERVER_PID=$! - sleep 5 - timeout 10 python mooncake-transfer-engine/tests/rpc_communicator_test.py client --url 127.0.0.1:9004 --threads 2 --data-size 1 || true - kill $SERVER_PID 2>/dev/null || true - wait $SERVER_PID 2>/dev/null || true - - - name: Test Safetensor Functions + - name: Run all CTest unit tests + id: ctest run: | - source test_env/bin/activate - pip install safetensors - python -m unittest mooncake-wheel.tests.test_safetensor_functions + reserved_ports=$(sysctl -n net.ipv4.ip_local_reserved_ports) + sudo sysctl -w "net.ipv4.ip_local_reserved_ports=${reserved_ports:+$reserved_ports,}50052" + mkdir -p build/test-results + 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 --test-dir build --parallel "$(nproc)" --output-on-failure \ + --output-junit "$GITHUB_WORKSPACE/build/test-results/ctest.xml" shell: bash + - name: Preserve CTest diagnostics + if: ${{ always() && steps.ctest.outcome == 'failure' }} + uses: ./.github/actions/ctest-diagnostics + with: + summary-title: PR CTest failure + artifact-name: ctest-diagnostics-pr + junit-report: build/test-results/ctest.xml + failed-tests-log: build/Testing/Temporary/LastTestsFailed.log + last-test-log: build/Testing/Temporary/LastTest.log + build-flags: - if: >- - github.event_name == 'push' || - github.event.action == 'opened' || - contains(github.event.pull_request.labels.*.name, 'run-ci') + needs: [check-paths] + if: *run-ci-for-source-changes runs-on: ubuntu-22.04 - strategy: - matrix: - python-version: ['3.10', '3.12'] env: CI: "true" - BUILD_WITH_EP: "1" - TORCH_CUDA_ARCH_LIST: "8.0;9.0" SCCACHE_GHA_ENABLED: "true" + PIP_NO_CACHE_DIR: "1" + CMAKE_RELWITHDEBINFO_FLAGS: "-O2 -DNDEBUG" steps: - uses: actions/checkout@v4 + with: + persist-credentials: false - - name: Set up Python ${{ matrix.python-version }} + - name: Set up Python 3.12 uses: actions/setup-python@v5 with: - python-version: ${{ matrix.python-version }} + python-version: '3.12' - name: Free up disk space - run: | - sudo rm -rf /usr/share/dotnet - sudo rm -rf /opt/ghc - sudo rm -rf /opt/hostedtoolcache/CodeQL - sudo rm -rf /usr/local/lib/android - df -h + uses: ./.github/actions/free-disk-space - name: Install CUDA Toolkit uses: Jimver/cuda-toolkit@v0.2.24 @@ -404,13 +494,10 @@ jobs: core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); - - name: Run sccache stat for check - shell: bash - run: ${SCCACHE_PATH} --show-stats - - name: Install dependencies run: | sudo apt update -y + sudo apt install -y ninja-build sudo bash -x dependencies.sh -y df -h shell: bash @@ -422,135 +509,98 @@ 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 .. -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 + 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" \ + "-DCMAKE_C_FLAGS_RELWITHDEBINFO=${CMAKE_RELWITHDEBINFO_FLAGS}" \ + "-DCMAKE_CXX_FLAGS_RELWITHDEBINFO=${CMAKE_RELWITHDEBINFO_FLAGS}" \ + -DENABLE_DEBUG_SYMBOLS=OFF + cmake --build . + sudo cmake --install . df -h - shell: bash - - - name: Configure project with all settings are ON - run: | - mkdir build - cd build - cmake .. -DUSE_ETCD=ON -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_STORE=ON -DWITH_P2P_STORE=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=ON -DUSE_MNNVL=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs" - shell: bash - # TODO: lack USE_NVMEOF,USE_MNNVL - - - name: Build project with all settings are ON - 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 - make -j4 - sudo make install + cd .. + rm -rf build df -h shell: bash - - name: Configure project with unit tests and examples - run: | - cd build - 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 - make -j4 - sudo make install - shell: bash - - - name: Configure project - run: | - cd build - rm -r */tests - cmake .. -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" - shell: bash - - - name: Build project - 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 - make -j4 - sudo make install - shell: bash + spell-check: + name: Spell Check with Typos + if: &run-ci >- + (github.event_name == 'push' || + github.event_name == 'workflow_dispatch' || + (github.event_name == 'pull_request' && + (github.event.action != 'labeled' || + 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 - - name: Build nvlink_allocator.so - run: | - mkdir -p build/mooncake-transfer-engine/nvlink-allocator - cd mooncake-transfer-engine/nvlink-allocator - export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH - export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH - export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH - bash build.sh ../../build/mooncake-transfer-engine/nvlink-allocator/ - shell: bash + reshard-type-check: + name: Check reshard manifest types + if: *run-ci + runs-on: ubuntu-22.04 + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false - - name: Generate Python version tag - id: generate_tag_flags - run: | - echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT - shell: bash + - name: Set up Python 3.12 + uses: actions/setup-python@v5 + with: + python-version: '3.12' - # In CI, build_wheel.sh removes build/ to free disk (CI=true); set FREE_BUILD_DIR=1 locally to enable. - - name: Build Python wheel + - name: Check canonical reshard contract types run: | - PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_flags.outputs.python_version_tag }} ./scripts/build_wheel.sh + python -m pip install --disable-pip-version-check pyright==1.1.411 + bash scripts/check_reshard_types.sh shell: bash - - name: Upload Python wheel artifact - uses: actions/upload-artifact@v4 - with: - name: mooncake-wheel-ubuntu-py${{ steps.generate_tag_flags.outputs.python_version_tag }} - path: mooncake-wheel/dist-py${{ steps.generate_tag_flags.outputs.python_version_tag }}/*.whl - - build-docker: - name: Build Docker Image - if: >- - github.event_name == 'push' || - github.event.action == 'opened' || - contains(github.event.pull_request.labels.*.name, 'run-ci') + release-gate-tests: + name: Test TestPyPI release gate + if: *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 + - name: Set up Python 3.12 + uses: actions/setup-python@v5 + with: + python-version: '3.12' - - name: Build Docker image + - name: Run focused release gate tests run: | - docker build -f docker/mooncake.Dockerfile \ - --build-arg PYTHON_VERSION=3.10 \ - --build-arg EP_TORCH_VERSIONS="2.9.1" \ - -t mooncake:from-source . - - spell-check: - name: Spell Check with Typos - if: >- - github.event_name == 'push' || - 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 - - name: Spell Check Repo - uses: crate-ci/typos@v1.30.2 + python -m pip install --disable-pip-version-check \ + packaging pytest==8.3.5 + python -m pytest -q mooncake-wheel/tests/test_testpypi_wheel_gate.py + shell: bash clang-format: name: Check code format - if: >- - github.event_name == 'push' || - github.event.action == 'opened' || - contains(github.event.pull_request.labels.*.name, 'run-ci') + if: *run-ci runs-on: ubuntu-22.04 steps: - name: Checkout Actions Repository uses: actions/checkout@v4 with: fetch-depth: 0 # Need full history for branch comparison + persist-credentials: false - name: Install clang-format 20 run: | @@ -566,7 +616,7 @@ jobs: echo "Error: code_format.sh not found or not executable" exit 1 fi - + # Determine base ref for comparison if [ "${{ github.event_name }}" == "pull_request" ]; then # For PRs: compare against the target branch @@ -583,5 +633,358 @@ jobs: fi fi echo "Comparing against: ${BASE_REF}" - ./scripts/code_format.sh --check --base "${BASE_REF}" + ./scripts/code_format.sh --check --changed-lines --base "${BASE_REF}" + shell: bash + + python-lint: + name: Check Python with Ruff + if: *run-ci + runs-on: ubuntu-22.04 + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + persist-credentials: false + + - name: Set up Python 3.12 + uses: actions/setup-python@v5 + with: + python-version: '3.12' + + - name: Install Ruff + run: python -m pip install --disable-pip-version-check ruff==0.6.9 + + - name: Check changed Python files + run: | + if [ "${{ github.event_name }}" = "pull_request" ]; then + base_ref="origin/${{ github.base_ref }}" + elif [ "${{ github.event.before }}" = \ + "0000000000000000000000000000000000000000" ]; then + base_ref="origin/${{ github.event.repository.default_branch }}" + elif [ -n "${{ github.event.before }}" ]; then + base_ref="${{ github.event.before }}" + else + base_ref="HEAD^" + fi + + mapfile -d '' python_files < <( + git diff --name-only --diff-filter=ACMR -z "$base_ref"...HEAD -- \ + '*.py' \ + ':(exclude)extern/**' \ + ':(exclude)FAST25-release/**' + ) + if [ "${#python_files[@]}" -eq 0 ]; then + echo "No changed Python files to check." + exit 0 + fi + + ruff check "${python_files[@]}" + ruff format --check "${python_files[@]}" + shell: bash + + + docs-check: + name: Check Sphinx docs build + if: *run-ci + runs-on: ubuntu-22.04 + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 2 + persist-credentials: false + + - uses: dorny/paths-filter@v3 + id: filter + with: + filters: | + docs: + - 'docs/**' + - 'docs/requirements-docs.txt' + + - name: Set up Python + if: steps.filter.outputs.docs == 'true' + uses: actions/setup-python@v5 + with: + python-version: '3.12' + + - name: Install documentation dependencies + if: steps.filter.outputs.docs == 'true' + run: | + python -m pip install --upgrade pip + pip install -r docs/requirements-docs.txt + + - name: Build docs with strict mode + if: steps.filter.outputs.docs == 'true' + run: | + cd docs + make html SPHINXOPTS=-W shell: bash + + + check-paths: + if: *run-ci + runs-on: ubuntu-latest + outputs: + should-run-downstream: ${{ steps.dispatch-override.outputs.src || steps.filter.outputs.src }} + should-run-tent: ${{ steps.dispatch-override.outputs.tent || steps.filter.outputs.tent }} + 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 + echo "tent=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/actions/**' + - '.github/workflows/**' + tent: + - 'mooncake-transfer-engine/**' + - 'mooncake-common/**' + - 'CMakeLists.txt' + - 'dependencies.sh' + - '.github/workflows/ci.yml' + + build-wheel-cu13: + needs: [spell-check, clang-format, check-paths] + if: *run-ci-for-source-changes + uses: ./.github/workflows/_build-wheel.yaml + with: + variant: cuda13 + python-versions: '["3.12"]' + artifact-prefix: mooncake-wheel-cu130 + version-override: 0.0.0.dev0 + + build-wheel-efa: + needs: [spell-check, clang-format, check-paths] + if: *run-ci-for-source-changes + uses: ./.github/workflows/ci_efa.yml + secrets: inherit + + build-wheel-rocm: + 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_name == 'pull_request' && + (github.event.action != 'labeled' || + contains(github.event.pull_request.labels.*.name, 'run-ci')))) + uses: ./.github/workflows/ci_rocm.yml + secrets: inherit + + tent-ci: + needs: [spell-check, clang-format, check-paths] + if: >- + (needs.check-paths.outputs.should-run-tent == 'true' || + github.event_name == 'workflow_dispatch') && + (github.event_name == 'push' || + github.event_name == 'workflow_dispatch' || + (github.event_name == 'pull_request' && + (github.event.action != 'labeled' || + contains(github.event.pull_request.labels.*.name, 'run-ci')))) + runs-on: ubuntu-22.04 + strategy: + fail-fast: false + matrix: + include: + - name: cuda-on + cmake_flags: '-DUSE_CUDA=ON -DCMAKE_EXE_LINKER_FLAGS=-L/usr/local/cuda/lib64/stubs' + need_cuda: true + metrics_flags: '' + - name: cuda-off + cmake_flags: '-DUSE_CUDA=OFF' + need_cuda: false + metrics_flags: '' + - name: cuda-off-metrics-on + cmake_flags: '-DUSE_CUDA=OFF' + need_cuda: false + metrics_flags: '-DTENT_METRICS_ENABLED=ON' + # Exercise the native UB target with its injected mock adapter. A + # Debug build also catches ODR/link errors hidden by optimization. + - name: ub-mock + cmake_flags: '-DUSE_CUDA=OFF -DUSE_UB=ON -DCMAKE_BUILD_TYPE=Debug' + need_cuda: false + metrics_flags: '' + name: tent-ci (${{ matrix.name }}) + env: + CI: "true" + SCCACHE_GHA_ENABLED: "true" + + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + + - name: Free up disk space + if: matrix.need_cuda + uses: ./.github/actions/free-disk-space + + - name: Install CUDA Toolkit + if: matrix.need_cuda + uses: Jimver/cuda-toolkit@v0.2.24 + with: + cuda: '12.8.1' + linux-local-args: '["--toolkit"]' + method: 'network' + sub-packages: '["nvcc", "nvrtc-dev"]' + + - name: Run sccache-cache + uses: mozilla-actions/sccache-action@v0.0.9 + + - name: Configure sccache + uses: actions/github-script@v7 + with: + script: | + core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); + core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); + + - 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: Configure project with TENT + run: | + mkdir build-tent + cd build-tent + cmake -G Ninja .. -DUSE_TENT=ON -DUSE_HTTP=ON -DENABLE_SCCACHE=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_DEBUG_SYMBOLS=OFF ${{ matrix.cmake_flags }} ${{ matrix.metrics_flags }} + shell: bash + + - name: Build project with TENT + run: | + if [ "${{ matrix.need_cuda }}" = "true" ]; then + export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH + export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH + fi + cd build-tent + cmake --build . + sudo cmake --install . + shell: bash + + # Check the Transfer Engine Rust crate against the freshly built C++ + # libraries. This is a compile/link check, so it is safe to run on every + # matrix leg (including the cuda-on leg, which only verifies that the + # USE_CUDA code paths compile). + - name: Check Transfer Engine Rust bindings, examples, and tests + run: | + if [ -f /usr/local/cuda/lib64/stubs/libcuda.so ] && \ + [ ! -e /usr/local/cuda/lib64/stubs/libcuda.so.1 ]; then + sudo ln -s libcuda.so /usr/local/cuda/lib64/stubs/libcuda.so.1 + fi + if [ "${{ matrix.need_cuda }}" = "true" ]; then + export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH + fi + cd mooncake-transfer-engine/rust + export MOONCAKE_BUILD_DIR=$GITHUB_WORKSPACE/build-tent + export MOONCAKE_TE_LIB_DIR=$GITHUB_WORKSPACE/build-tent/mooncake-transfer-engine/src + export MOONCAKE_TE_INCLUDE_DIR=$GITHUB_WORKSPACE/mooncake-transfer-engine/include + cargo test --lib + cargo test --examples --tests --no-run + cargo clean + shell: bash + + # Only run tests on the cuda-off leg. GitHub runners have no real GPU; + # with USE_CUDA=ON tent's cuda_probe hits the CUDA stub library at + # runtime and drives some dispatch paths past the fake objects the + # unit tests rely on, causing false failures. cuda-on still validates + # that every #ifdef USE_CUDA branch compiles. + - name: Test (TENT) + id: ctest + if: '!matrix.need_cuda' + run: | + cd build-tent + mkdir -p test-results + ctest --test-dir mooncake-transfer-engine/tent/tests -j \ + --output-on-failure \ + --output-junit "$GITHUB_WORKSPACE/build-tent/test-results/ctest.xml" + shell: bash + + - name: Preserve TENT CTest diagnostics + if: ${{ always() && steps.ctest.outcome == 'failure' }} + uses: ./.github/actions/ctest-diagnostics + with: + summary-title: TENT CTest failure (${{ matrix.name }}) + artifact-name: ctest-diagnostics-tent-${{ matrix.name }} + junit-report: build-tent/test-results/ctest.xml + failed-tests-log: build-tent/mooncake-transfer-engine/tent/tests/Testing/Temporary/LastTestsFailed.log + last-test-log: build-tent/mooncake-transfer-engine/tent/tests/Testing/Temporary/LastTest.log + + - name: Start Metadata Server + if: '!matrix.need_cuda' + run: | + cd mooncake-transfer-engine/example/http-metadata-server-python + pip install aiohttp + python ./bootstrap_server.py & + sleep 2 + shell: bash + + - name: Run Transfer Engine Rust smoke test + if: '!matrix.need_cuda' + env: + MOONCAKE_BUILD_DIR: ${{ github.workspace }}/build-tent + MOONCAKE_TE_RUST_LINK_ASAN: "0" + run: ./scripts/ci/run_transfer_engine_rust_smoke.sh + shell: bash + + - name: Run sccache stat for check + if: ${{ env.SCCACHE_PATH != '' }} + shell: bash + run: ${SCCACHE_PATH} --show-stats + + ci-gate: + name: CI Gate + if: always() + needs: + - spell-check + - clang-format + - python-lint + - docs-check + - reshard-type-check + - release-gate-tests + - build-wheel + - unit-tests + - build-flags + - test-wheel-ubuntu + - build-wheel-cu13 + - build-wheel-efa + - build-wheel-rocm + - tent-ci + 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) }} diff --git a/.github/workflows/ci_ascend.yml b/.github/workflows/ci_ascend.yml index 1d961b988a..1f2b6d1ea2 100644 --- a/.github/workflows/ci_ascend.yml +++ b/.github/workflows/ci_ascend.yml @@ -1,11 +1,12 @@ name: 'CI Test on ASCEND Platform' on: - push: - branches: [ "main" ] - pull_request: - branches: [ "main" ] - types: [opened, synchronize, reopened, labeled] + workflow_call: + inputs: + checkout_ref: + description: 'Git ref to checkout (PR head SHA)' + required: false + type: string jobs: build-and-test: @@ -14,8 +15,8 @@ jobs: container: image: localhost:5000/mooncake-hixl-ci:v5 - options: --privileged --user 0:0 --device /dev/davinci0 --device /dev/davinci1 --device /dev/davinci2 --device /dev/davinci3 - --device /dev/davinci4 --device /dev/davinci5 --device /dev/davinci6 --device /dev/davinci7 + options: --privileged --user 0:0 --device /dev/davinci0 --device /dev/davinci1 --device /dev/davinci2 --device /dev/davinci3 + --device /dev/davinci4 --device /dev/davinci5 --device /dev/davinci6 --device /dev/davinci7 --device /dev/davinci_manager --device /dev/devmm_svm --device /dev/hisi_hdc --ulimit nproc=65535:65535 env: GITHUB_ACTIONS: "true" @@ -27,30 +28,128 @@ jobs: - /etc/hccn.conf:/etc/hccn.conf steps: + - name: Configure GitHub fetch defaults + shell: bash + run: | + git config --global protocol.version 2 + git config --global http.version HTTP/1.1 + git config --global http.lowSpeedLimit 1024 + git config --global http.lowSpeedTime 30 + - name: Checkout code + id: checkout_code + continue-on-error: true + uses: actions/checkout@v4 + with: + ref: ${{ inputs.checkout_ref || github.sha }} + fetch-depth: 1 + persist-credentials: false + + # Mirror is only a git transport fallback. Keep retry on actions/checkout so + # pull_request_target fork checks still run; do not fetch with raw git. + - name: Configure GitHub mirror rewrite + if: steps.checkout_code.outcome == 'failure' + shell: bash + env: + ASCEND_GITHUB_MIRROR_URLS: 'https://ghfast.top/' + 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" + find "$workdir" -mindepth 1 -maxdepth 1 -exec rm -rf {} + + + # insteadOf only rewrites github.com fetches; checkout still runs + # assertSafePrCheckout before any git network I/O. + mirror_base="${candidates[0]}" + echo "Rewriting https://github.com/ to ${mirror_base}https://github.com/" + git config --global url."${mirror_base}https://github.com/".insteadOf "https://github.com/" + + - name: Retry checkout via GitHub mirror + if: steps.checkout_code.outcome == 'failure' uses: actions/checkout@v4 with: + ref: ${{ inputs.checkout_ref || github.sha }} fetch-depth: 1 persist-credentials: false - name: Configure CMake shell: bash + env: + ASCEND_GITHUB_MIRROR_URLS: 'https://ghfast.top/' run: | source /usr/local/Ascend/cann-9.0.0/set_env.sh pwd - if ! git submodule update --init --recursive; then - if [ ! -d "extern/pybind11" ] || [ -z "$(ls -A 'extern/pybind11' 2>/dev/null)" ]; then - echo "git submodule update failed, try to cp pybind11..." - if [ -d "../pybind11" ]; then - cp -r ../pybind11 extern/ - else - echo "Error: ../pybind11 does not exist. Cannot copy pybind11." - exit 1 - fi + + 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 "Detected that extern/pybind11 already exists, continuing execution...." + 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 @@ -59,7 +158,9 @@ jobs: cmake .. \ -DUSE_ASCEND_DIRECT=ON \ -DBUILD_EXAMPLES=OFF \ - -DBUILD_UNIT_TESTS=OFF + -DBUILD_UNIT_TESTS=OFF \ + -DENABLE_DEBUG_SYMBOLS=OFF \ + -DBUILD_BENCHMARK=OFF - name: Build shell: bash @@ -67,8 +168,8 @@ jobs: source /usr/local/Ascend/cann-9.0.0/set_env.sh echo "Building..." cd build - make -j8 - make install -j8 + cmake --build . -j48 + cmake --install . echo "Mooncake installed successfully." - name: Run Hixl Mooncake Store Test @@ -117,7 +218,7 @@ jobs: # Check if master is running if ! kill -0 $MASTER_PID 2>/dev/null; then echo "Error: Mooncake Master failed to start" - cat /tmp/mooncake_master.log + cat /tmp/mooncake_master.log exit 1 fi @@ -129,37 +230,86 @@ jobs: "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 case - export ASCEND_BUFFER_POOL=4:8 - for test_case in "${TEST_CASES[@]}"; do + # Run each test scenario + for scenario in "${TEST_SCENARIOS[@]}"; do echo "" echo "=========================================" - echo "Running test: $test_case" + echo "Running scenario: $scenario" echo "=========================================" - if [ ! -f "$test_case" ]; then - echo "Warning: Test file $test_case not found, skipping..." - continue + # 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 the test with run.sh - python3 $test_case \ - --device_id=0 \ - --rank=0 \ - 2>&1 | tee "/tmp/hixl_test_${test_case%.py}.log" - TEST_RESULT=${PIPESTATUS[0]} + # Run each test case in the current scenario + for test_case in "${TEST_CASES[@]}"; do + echo "" + echo "-----------------------------------------" + echo "Test: $test_case" + echo "-----------------------------------------" - if [ $TEST_RESULT -eq 0 ]; then - echo "✓ $test_case PASSED" - PASSED_TESTS+=("$test_case") - else - echo "✗ $test_case FAILED with exit code: $TEST_RESULT" - FAILED_TESTS+=("$test_case") - fi + 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 "" @@ -192,7 +342,7 @@ jobs: echo "" echo "All Hixl Mooncake Store tests completed successfully!" - + - name: Test Summary if: always() diff --git a/.github/workflows/ci_cu13.yml b/.github/workflows/ci_cu13.yml deleted file mode 100644 index c51016ebad..0000000000 --- a/.github/workflows/ci_cu13.yml +++ /dev/null @@ -1,129 +0,0 @@ -name: 'Build Wheel (CUDA 13)' - -on: - 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: - python-version: ['3.10', '3.12'] - env: - BUILD_WITH_EP: "1" - CU13_BUILD: "1" - TORCH_CUDA_ARCH_LIST: "8.0;9.0" - SCCACHE_GHA_ENABLED: "true" - - steps: - - uses: actions/checkout@v4 - - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - - - name: Free up disk space - run: | - sudo rm -rf /usr/share/dotnet - sudo rm -rf /opt/ghc - sudo rm -rf /opt/hostedtoolcache/CodeQL - sudo rm -rf /usr/local/lib/android - df -h - - - name: Install CUDA Toolkit - uses: Jimver/cuda-toolkit@v0.2.29 - with: - cuda: '13.0.2' - linux-local-args: '["--toolkit"]' - method: 'network' - sub-packages: '["nvcc", "nvrtc-dev"]' - non-cuda-sub-packages: '["libcusparse-dev", "libcublas-dev", "libcusolver-dev"]' - - - name: Run sccache-cache - uses: mozilla-actions/sccache-action@v0.0.9 - - - name: Configure sccache - uses: actions/github-script@v7 - with: - script: | - core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); - core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); - - - name: Run sccache stat for check - shell: bash - run: ${SCCACHE_PATH} --show-stats - - - name: Install dependencies - run: | - sudo apt update -y - sudo bash -x dependencies.sh -y - df -h - shell: bash - - - name: Configure project - run: | - mkdir build - cd build - 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" \ - -DWITH_METRICS=ON \ - -DBUILD_UNIT_TESTS=OFF \ - -DBUILD_EXAMPLES=ON \ - -DENABLE_SCCACHE=ON \ - -DBUILD_BENCHMARK=ON \ - -DUSE_CUDA=ON \ - -DCMAKE_BUILD_TYPE=Release \ - -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs" - shell: bash - - - name: Build project - 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 - make -j4 - sudo make install - df -h - shell: bash - - - name: Build nvlink_allocator.so - run: | - mkdir -p build/mooncake-transfer-engine/nvlink-allocator - cd mooncake-transfer-engine/nvlink-allocator - export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH - export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH - export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH - bash build.sh ../../build/mooncake-transfer-engine/nvlink-allocator/ - shell: bash - - - name: Generate Python version tag - id: generate_tag - run: | - echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT - shell: bash - - - name: Build Python wheel - run: | - PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag.outputs.python_version_tag }} ./scripts/build_wheel.sh - shell: bash - - - name: Upload Python wheel artifact - uses: actions/upload-artifact@v4 - with: - name: mooncake-wheel-cu130-ubuntu-py${{ steps.generate_tag.outputs.python_version_tag }} - path: mooncake-wheel/dist-py${{ steps.generate_tag.outputs.python_version_tag }}/*.whl diff --git a/.github/workflows/ci_efa.yml b/.github/workflows/ci_efa.yml new file mode 100644 index 0000000000..6bf8967c72 --- /dev/null +++ b/.github/workflows/ci_efa.yml @@ -0,0 +1,41 @@ +name: 'Build Wheel (AWS EFA)' + +on: + workflow_call: {} + +# No EFA hardware is required to build. PR validation covers the CUDA and +# non-CUDA variants with one Python version each; releases use the full matrix. +jobs: + build-wheel-efa: + strategy: + matrix: + include: + - variant: cuda + use_cuda: "ON" + build_env: "EFA_BUILD" + cuda-version: "12.8.1" + python-version: "3.12" + - variant: cuda13 + use_cuda: "ON" + build_env: "EFA_CU13_BUILD" + cuda-version: "13.0.2" + python-version: "3.12" + - variant: non-cuda + use_cuda: "OFF" + build_env: "EFA_NON_CUDA_BUILD" + cuda-version: "12.8.1" + python-version: "3.10" + uses: ./.github/workflows/_build-efa-wheel.yaml + with: + variant: ${{ matrix.variant }} + use-cuda: ${{ matrix.use_cuda == 'ON' }} + python-versions: ${{ format('["{0}"]', matrix.python-version) }} + build-profile: ci + cmake-args: >- + -DUSE_ETCD=ON -DUSE_HTTP=ON -DWITH_STORE=ON -DWITH_METRICS=ON + -DBUILD_UNIT_TESTS=OFF -DBUILD_EXAMPLES=ON -DBUILD_BENCHMARK=ON + -DENABLE_DEBUG_SYMBOLS=OFF + variant-flag: ${{ matrix.build_env }} + cuda-version: ${{ matrix.cuda-version }} + torch-cuda-arch-list: '8.0;9.0' + artifact-prefix: mooncake-wheel-efa-${{ matrix.variant }}-ubuntu diff --git a/.github/workflows/ci_rocm.yml b/.github/workflows/ci_rocm.yml new file mode 100644 index 0000000000..c445be1230 --- /dev/null +++ b/.github/workflows/ci_rocm.yml @@ -0,0 +1,157 @@ +name: 'Build Wheel (ROCm)' + +# ROCm/HIP CI parity with the standard CUDA wheel CI: build the AMD ROCm wheel +# on PRs so packaging regressions are caught. Runs inside the ROCm dev image so +# hipcc / HIP headers / hsa-runtime are available; no GPU is needed to compile. +# On-device transfer tests (test_transfer_on_hip.py) require AMD hardware and +# run outside GitHub-hosted runners. + +on: + workflow_call: {} + +jobs: + build-wheel-rocm: + runs-on: ubuntu-22.04 + # Pinned to match the upstream vllm/vllm-openai-rocm and sglang ROCm images + # (ROCm 7.2). The wheel excludes the ROCm runtime and binds it at load time, + # so it also loads on the ROCm 7.0 image variants. + container: rocm/dev-ubuntu-22.04:7.2.3-complete + strategy: + matrix: + # 3.10 covers the SGLang ROCm image; 3.12 covers vLLM ROCm. + python-version: ['3.10', '3.12'] + env: + HIP_BUILD: "1" + SCCACHE_GHA_ENABLED: "true" + + steps: + # git must exist BEFORE checkout so actions/checkout does a real clone + # (with .git + submodules). The rocm/dev image ships without git, which + # otherwise makes checkout fall back to a source tarball and breaks + # dependencies.sh's `git submodule` step. + - name: Install git (pre-checkout) + shell: bash + run: | + set -eo pipefail + export DEBIAN_FRONTEND=noninteractive + apt-get update -y + apt-get install -y --no-install-recommends git ca-certificates + + - uses: actions/checkout@v4 + with: + persist-credentials: false + submodules: recursive + + - name: Mark repository as safe + run: git config --global --add safe.directory "$GITHUB_WORKSPACE" || true + shell: bash + + - name: Install toolchain and Python ${{ matrix.python-version }} + shell: bash + run: | + set -eo pipefail + export DEBIAN_FRONTEND=noninteractive + apt-get update -y + apt-get install -y --no-install-recommends \ + curl build-essential sudo pkg-config \ + ninja-build software-properties-common + PYV="${{ matrix.python-version }}" + if ! command -v "python${PYV}" >/dev/null 2>&1; then + add-apt-repository -y ppa:deadsnakes/ppa + apt-get update -y + fi + # Always install -dev + -venv: the image's system python3.10 exists but + # ships without the venv module / dev headers. + apt-get install -y --no-install-recommends \ + "python${PYV}" "python${PYV}-dev" "python${PYV}-venv" + curl -sS https://bootstrap.pypa.io/get-pip.py | "python${PYV}" + PYTHON_BIN="$(command -v python${PYV})" + echo "PYTHON_BIN=${PYTHON_BIN}" >> "$GITHUB_ENV" + # hipify-perl lives in /opt/rocm/bin; CMake's USE_HIP path requires it. + # Use GITHUB_PATH (setting PATH via GITHUB_ENV is not honored reliably). + echo "/opt/rocm/bin" >> "$GITHUB_PATH" + + - name: Run sccache-cache + uses: mozilla-actions/sccache-action@v0.0.9 + + - name: Configure sccache + uses: actions/github-script@v7 + with: + script: | + core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); + core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); + + - name: Install dependencies + shell: bash + run: | + set -eo pipefail + bash -x dependencies.sh -y + echo "/usr/local/go/bin" >> "$GITHUB_PATH" + + - name: Configure project + shell: bash + run: | + set -eo pipefail + export PATH="/opt/rocm/bin:$PATH" # hipify-perl for USE_HIP + rm -rf build && mkdir build && cd build + cmake -G Ninja .. \ + -DUSE_HIP=ON -DUSE_CUDA=OFF -DWITH_EP=OFF \ + -DUSE_HTTP=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON \ + -DBUILD_UNIT_TESTS=OFF -DENABLE_SCCACHE=ON \ + -DCMAKE_BUILD_TYPE=Release \ + -DPython3_EXECUTABLE="${PYTHON_BIN}" + + - name: Build project + shell: bash + run: | + set -eo pipefail + export PATH="/opt/rocm/bin:$PATH" + cd build + # Retry to ride out transient Go module (proxy.golang.org) fetch errors + # during the etcd-wrapper build; Go caches modules, so retries resume. + n=0 + until cmake --build . -j"$(nproc)"; do + n=$((n+1)) + if [ "$n" -ge 3 ]; then echo "Build failed after $n attempts"; exit 1; fi + echo "Build attempt $n failed; retrying in 15s..."; sleep 15 + done + cmake --install . + + - name: Run sccache stat for check + if: ${{ env.SCCACHE_PATH != '' }} + shell: bash + run: ${SCCACHE_PATH} --show-stats + + - name: Generate Python version tag + id: generate_tag + shell: bash + run: echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> "$GITHUB_OUTPUT" + + - name: Build Python wheel + shell: bash + run: | + set -eo pipefail + export PATH="/opt/rocm/bin:$PATH" + export LD_LIBRARY_PATH="${LD_LIBRARY_PATH:-}:/usr/local/lib" + HIP_BUILD=1 PYTHON_VERSION=${{ matrix.python-version }} \ + OUTPUT_DIR=dist-rocm-py${{ steps.generate_tag.outputs.python_version_tag }} \ + ./scripts/build_wheel.sh + + - name: Smoke test wheel + shell: bash + run: | + set -eo pipefail + smoke_venv=$(mktemp -d) + "${PYTHON_BIN}" -m venv "$smoke_venv" + "$smoke_venv/bin/python" -m pip install --no-deps \ + mooncake-wheel/dist-rocm-py${{ steps.generate_tag.outputs.python_version_tag }}/*.whl + # ROCm runtime is excluded from the wheel and bound at load time. + export LD_LIBRARY_PATH="/opt/rocm/lib:/usr/local/lib:${LD_LIBRARY_PATH:-}" + site="$("$smoke_venv/bin/python" -c 'import mooncake,os;print(os.path.dirname(mooncake.__file__))')" + "$site/mooncake_master" --version + + - name: Upload Python wheel artifact + uses: actions/upload-artifact@v4 + with: + name: mooncake-wheel-rocm-ubuntu-py${{ steps.generate_tag.outputs.python_version_tag }} + path: mooncake-wheel/dist-rocm-py${{ steps.generate_tag.outputs.python_version_tag }}/*.whl diff --git a/.github/workflows/deploy.yml b/.github/workflows/deploy.yml index 0456939148..6862480bae 100644 --- a/.github/workflows/deploy.yml +++ b/.github/workflows/deploy.yml @@ -4,6 +4,10 @@ on: # Runs on pushes targeting the default branch push: branches: ["main"] + paths: + - 'docs/**' + - 'docs/requirements-docs.txt' + - '.github/workflows/deploy.yml' # Allows you to run this workflow manually from the Actions tab workflow_dispatch: @@ -31,31 +35,33 @@ jobs: steps: - name: Checkout uses: actions/checkout@v4 - + with: + persist-credentials: false + - name: Setup Python uses: actions/setup-python@v4 with: python-version: '3.x' # Choose the specific version as needed - + - name: Install dependencies run: | - pip install -r requirements_docs.txt - + pip install -r docs/requirements-docs.txt + - name: Build documentation run: | cd docs make clean make html - + - name: Setup Pages uses: actions/configure-pages@v5 - + - name: Upload artifact uses: actions/upload-pages-artifact@v3 with: # Upload only the build directory path: './docs/build/html' - + - name: Deploy to GitHub Pages id: deployment - uses: actions/deploy-pages@v4 \ No newline at end of file + uses: actions/deploy-pages@v4 diff --git a/.github/workflows/e2e-ci.yml b/.github/workflows/e2e-ci.yml new file mode 100644 index 0000000000..b6abfee2e8 --- /dev/null +++ b/.github/workflows/e2e-ci.yml @@ -0,0 +1,94 @@ +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 diff --git a/.github/workflows/integration-test.yml b/.github/workflows/integration-test.yml index 92dee2ef7b..3888e68237 100644 --- a/.github/workflows/integration-test.yml +++ b/.github/workflows/integration-test.yml @@ -1,23 +1,29 @@ name: 'Integration test (Linux)' on: - push: - branches: [ "main" ] - paths: - - 'mooncake-*/**' - - 'extern/**' - - 'CMakeLists.txt' - - 'scripts/**' - - '.github/workflows/**' - pull_request_target: - branches: [ "main" ] - types: [opened, synchronize, reopened, labeled] - paths: - - 'mooncake-*/**' - - 'extern/**' - - 'CMakeLists.txt' - - 'scripts/**' - - '.github/workflows/**' + 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 + source_run_id: + description: 'Workflow run containing the wheel artifact' + required: false + type: string + artifact_name: + description: 'Exact wheel artifact name' + required: false + type: string + testpypi_version: + description: 'Exact TestPyPI CUDA 13 package version to test instead of a CI artifact' + required: false + default: '' + type: string jobs: test-sglang-integration: @@ -25,41 +31,117 @@ jobs: env: tone_user_name: ${{ secrets.TONE_USER_NAME }} steps: + - name: Require T-one credentials for gated integration + if: ${{ inputs.testpypi_version != '' || inputs.source_run_id != '' || inputs.artifact_name != '' }} + env: + TONE_USER_NAME: ${{ secrets.TONE_USER_NAME }} + TONE_USER_TOKEN: ${{ secrets.TONE_USER_TOKEN }} + run: | + if [ -z "$TONE_USER_NAME" ] || [ -z "$TONE_USER_TOKEN" ]; then + echo "TONE_USER_NAME and TONE_USER_TOKEN are required for gated integration" + exit 1 + fi + + - name: Set up Python for TestPyPI download + if: ${{ inputs.testpypi_version != '' }} + uses: actions/setup-python@v5 + with: + python-version: '3.12' + + - name: Download the exact indexed CUDA 13 wheel + if: ${{ inputs.testpypi_version != '' }} + env: + TESTPYPI_VERSION: ${{ inputs.testpypi_version }} + run: | + python -m pip download \ + --index-url https://test.pypi.org/simple \ + --no-cache-dir \ + --no-deps \ + --only-binary=:all: \ + --pre \ + --dest integration-wheel \ + "mooncake-transfer-engine-cuda13==${TESTPYPI_VERSION}" + + - name: Stage the indexed wheel for T-one + if: ${{ inputs.testpypi_version != '' }} + id: testpypi-wheel + uses: actions/upload-artifact@v4 + with: + name: mooncake-testpypi-tone-${{ inputs.testpypi_version }} + path: integration-wheel/*.whl + if-no-files-found: error + retention-days: 14 + - name: trigger T-one test if: ${{ env.tone_user_name != '' }} + env: + TESTPYPI_ARTIFACT_ID: ${{ steps.testpypi-wheel.outputs.artifact-id }} run: | - SHA="${{ github.event.pull_request.head.sha }}" - PR_ID="${{ github.event.pull_request.number }}" + # 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 }}" + SOURCE_RUN_ID="${{ inputs.source_run_id }}" + ARTIFACT_NAME="${{ inputs.artifact_name }}" if [ "${{ github.event_name }}" = "push" ]; then SHA="${{ github.sha }}" PR_ID="" fi echo "PR_ID=${PR_ID}" + artifact_id="$TESTPYPI_ARTIFACT_ID" max_attempts=120 attempt=1 - while [ $attempt -le $max_attempts ]; do + while [ -z "$artifact_id" ] && [ $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 - artifact_id="" - if jq empty artifact.json >/dev/null 2>&1; then - 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) + echo "Target SHA=${SHA}" + artifact_id="" + run_id="$SOURCE_RUN_ID" + if [ -z "$run_id" ]; then + if curl -L -fs -o runs.json -H "Accept: application/vnd.github+json" -H "X-GitHub-Api-Version: 2022-11-28" "https://api.github.com/repos/${{ github.repository }}/actions/runs?head_sha=${SHA}&per_page=100"; then + if jq empty runs.json >/dev/null 2>&1; then + run_id=$(jq -r '.workflow_runs[] | select((.path == ".github/workflows/ci.yml") or (.name == "Build & Test (Linux)")) | .id' runs.json | head -n 1) + else + echo "Failed to download workflow run list. Retrying..." + fi else - echo "Failed to download artifact list. Retrying..." + echo "Failed to fetch workflow runs. Retrying..." fi - if [ -n "$artifact_id" ]; then - echo "Successfully fetched expected artifact id $artifact_id" - break - else - echo "Failed to fetch expected artifact. Retrying..." - if [ $attempt -lt $max_attempts ]; then - sleep $((attempt * 60 < 600 ? attempt * 60 : 600)) + fi + + if [ -n "$run_id" ]; then + echo "Matched workflow run id $run_id" + 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/runs/${run_id}/artifacts?per_page=100"; then + if jq empty artifact.json >/dev/null 2>&1; then + if [ -n "$ARTIFACT_NAME" ]; then + artifact_id=$(jq -r --arg name "$ARTIFACT_NAME" '.artifacts[] | select(.name == $name) | .id' artifact.json | head -n 1) + else + artifact_id=$(jq -r '.artifacts[] | select(.name | contains("py312") ) | select(.name | contains("mooncake") ) | select(.name | contains("cu130") ) | .id' artifact.json | head -n 1) + fi + if [ -z "$artifact_id" ]; then + echo "Available artifacts in workflow run $run_id:" + jq -r '.artifacts[].name' artifact.json || true + fi + else + echo "Failed to download artifact list. Retrying..." fi + else + echo "Failed to fetch artifacts for workflow run $run_id. Retrying..." fi else - echo "Failed to fetch artifacts. Retrying..." + echo "Failed to find Build & Test workflow run for SHA $SHA. Retrying..." + if jq empty runs.json >/dev/null 2>&1; then + echo "Available workflow runs for SHA:" + jq -r '.workflow_runs[] | "\(.id) \(.name) \(.path) \(.status) \(.conclusion)"' runs.json || true + fi + fi + + if [ -n "$artifact_id" ]; then + echo "Successfully fetched expected artifact id $artifact_id" + break + else + echo "Failed to fetch expected artifact. Retrying..." if [ $attempt -lt $max_attempts ]; then - sleep $((attempt * 60)) + sleep $((attempt * 60 < 600 ? attempt * 60 : 600)) fi fi attempt=$((attempt + 1)) @@ -70,6 +152,9 @@ jobs: fi ENV_INFO="ARTIFACT_ID=${artifact_id} GIT_REPO=${{ github.repository }}" + if [ -n "${{ inputs.testpypi_version }}" ]; then + ENV_INFO="${ENV_INFO} BRANCH=${{ github.ref_name }}" + fi if [ -n "$PR_ID" ]; then ENV_INFO="${ENV_INFO} PR_ID=${PR_ID}" fi diff --git a/.github/workflows/issue-bot.yml b/.github/workflows/issue-bot.yml new file mode 100644 index 0000000000..c199c04996 --- /dev/null +++ b/.github/workflows/issue-bot.yml @@ -0,0 +1,71 @@ +name: Issue Bot + +on: + issues: + types: [opened] + schedule: + # Daily at 03:00 UTC + - cron: "0 3 * * *" + workflow_dispatch: + +permissions: + issues: write + +jobs: + welcome-and-assign: + if: github.event_name == 'issues' + runs-on: ubuntu-latest + steps: + - name: Post welcome comment + uses: actions/github-script@v7 + with: + script: | + const issue = context.payload.issue; + const body = [ + `Thanks for opening this issue, @${issue.user.login}!`, + '', + '| Field | Value |', + '|-------|-------|', + `| **Issue** | #${issue.number} |`, + `| **GitHub user ID** | \`${issue.user.id}\` |`, + `| **Reporter** | @${issue.user.login} |`, + '', + 'A maintainer will triage this when possible. To help us respond faster, please include:', + '', + '- Mooncake version or commit SHA', + '- Environment (OS, CUDA/driver, RDMA stack if relevant)', + '- Steps to reproduce and expected vs. actual behavior', + '', + `Useful links: [Documentation](https://kvcache-ai.github.io/Mooncake/) · [Contributing guide](https://github.com/${context.repo.owner}/${context.repo.repo}/blob/main/CONTRIBUTING.md)`, + '', + '> This message was posted automatically by the issue bot.', + ].join('\n'); + + await github.rest.issues.createComment({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: issue.number, + body, + }); + + close-stale-issues: + if: github.event_name != 'issues' + runs-on: ubuntu-latest + steps: + - name: Close inactive issues (3+ months) + uses: actions/stale@v9 + with: + repo-token: ${{ secrets.GITHUB_TOKEN }} + days-before-stale: 90 + days-before-close: 7 + stale-issue-label: stale + close-issue-label: auto-closed + stale-issue-message: > + This issue has had no activity for 90 days and will be closed in 7 days + if there is no further activity. Please comment or react if it should + stay open. + close-issue-message: > + Closing due to 3 months of inactivity. If this is still relevant, + please comment and we can reopen. + exempt-issue-labels: pinned,keep-open,security + operations-per-run: 100 diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml new file mode 100644 index 0000000000..55a4675c32 --- /dev/null +++ b/.github/workflows/nightly.yml @@ -0,0 +1,713 @@ +name: Nightly Build & Test + +on: + schedule: + - cron: '0 16 * * *' # 00:00 Beijing Time (UTC+8) + workflow_dispatch: + inputs: + skip_publish: + description: 'Skip publishing to TestPyPI' + type: boolean + default: false + +permissions: + contents: read + +concurrency: + group: nightly-${{ github.ref_name }} + cancel-in-progress: true + +# ── Workflow Design Notes ────────────────────────────────────────────────────── +# 1. version-stamp computes a forward-looking nightly version (next patch + .devYYYYMMDD) +# that sorts AFTER the current release per PEP 440, e.g. "0.3.12.dev20260720". +# 2. Build jobs pass the nightly version via `version-override` to _build-wheel.yaml +# so each wheel is stamped with the correct dev version for TestPyPI. +# 3. Specialized platform and packaging checks run nightly instead of on every PR. +# 4. publish-testpypi uploads all nightly-* artifacts to TestPyPI. +# 5. nightly-gate aggregates ALL build + test results; notify-failure files an issue. +# ────────────────────────────────────────────────────────────────────────────── + +env: + SCCACHE_GHA_ENABLED: "true" + +# Gate the dependency root so all downstream jobs are skipped in forks. +jobs: + version-stamp: + if: github.repository == 'kvcache-ai/Mooncake' + runs-on: ubuntu-22.04 + outputs: + nightly_version: ${{ steps.version.outputs.nightly_version }} + steps: + - uses: actions/checkout@v4 + - name: Compute nightly version + id: version + run: | + BASE_VERSION=$(grep -Po '(?<=^version = ")[^"]+' mooncake-wheel/pyproject.toml) + # Strip .postN suffix if present, then bump patch for forward-looking version + CLEAN_VERSION=$(echo "$BASE_VERSION" | sed 's/\.post[0-9]*$//') + MAJOR=$(echo "$CLEAN_VERSION" | cut -d. -f1) + MINOR=$(echo "$CLEAN_VERSION" | cut -d. -f2) + PATCH=$(echo "$CLEAN_VERSION" | cut -d. -f3) + NEXT_PATCH=$((PATCH + 1)) + NIGHTLY_VERSION="${MAJOR}.${MINOR}.${NEXT_PATCH}.dev$(date -u +%Y%m%d)" + echo "nightly_version=$NIGHTLY_VERSION" >> "$GITHUB_OUTPUT" + echo "Nightly version: $NIGHTLY_VERSION" + + build-wheels: + needs: version-stamp + strategy: + fail-fast: false + matrix: + include: + - variant: cuda + architecture: x86_64 + artifact-prefix: nightly-cuda128-x86 + - variant: cuda + architecture: arm64 + artifact-prefix: nightly-cuda128-arm64 + - variant: cuda13 + architecture: x86_64 + artifact-prefix: nightly-cuda13-x86 + - variant: cuda13 + architecture: arm64 + artifact-prefix: nightly-cuda13-arm64 + - variant: non-cuda + architecture: x86_64 + artifact-prefix: nightly-non-cuda-x86 + - variant: non-cuda + architecture: arm64 + artifact-prefix: nightly-non-cuda-arm64 + uses: ./.github/workflows/_build-wheel.yaml + with: + variant: ${{ matrix.variant }} + architecture: ${{ matrix.architecture }} + artifact-prefix: ${{ matrix.artifact-prefix }} + version-override: ${{ needs.version-stamp.outputs.nightly_version }} + + tone-sglang-integration: + needs: build-wheels + uses: ./.github/workflows/integration-test.yml + with: + pr_sha: ${{ github.sha }} + source_run_id: ${{ github.run_id }} + artifact_name: nightly-cuda13-x86-py312 + secrets: inherit + + build-musa: + runs-on: ubuntu-22.04 + container: registry.mthreads.com/mcconline/inference/pytorch:2.9.1.post1-py3.10-musa5.2.0-mp31-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 + shell: bash + + - 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 -DWITH_EP=ON -DBUILD_UNIT_TESTS=OFF -DBUILD_EXAMPLES=OFF -DENABLE_DEBUG_SYMBOLS=OFF + shell: bash + + - name: Build project + run: | + cd build + source ~/.bashrc + # Limit concurrency of Mooncake PG's setup.py build to avoid OOM + export MAX_JOBS=2 + # Does not affect the Ninja system + cmake --build . + cmake --install . + shell: bash + + build-docker: + name: Build Docker Image + runs-on: ubuntu-22.04 + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + + - name: Free up disk space + uses: ./.github/actions/free-disk-space + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Build Docker image + run: | + docker build -f docker/mooncake.Dockerfile \ + --build-arg PYTHON_VERSION=3.10 \ + --build-arg EP_TORCH_VERSIONS="2.13.0" \ + --build-arg CLEAN_BUILD_ARTIFACTS=1 \ + -t mooncake:from-source . + + ascend-test: + uses: ./.github/workflows/ci_ascend.yml + secrets: inherit + + publish-testpypi: + needs: [version-stamp, build-wheels, nightly-test] + if: ${{ !(github.event_name == 'workflow_dispatch' && github.event.inputs.skip_publish == 'true') }} + runs-on: ubuntu-22.04 + environment: nightly + permissions: + contents: read + steps: + - name: Log nightly version + run: | + echo "Publishing nightly version: ${{ needs.version-stamp.outputs.nightly_version }}" + + - name: Download all nightly wheel artifacts + uses: actions/download-artifact@v4 + with: + path: dist-all + pattern: 'nightly-*' + + - name: Collect and rename wheels with nightly version + run: | + mkdir -p dist-publish + find dist-all -name "*.whl" -exec cp {} dist-publish/ \; + echo "Collected wheels:" + ls -la dist-publish/ + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.12' + + - name: Install twine + run: pip install twine + + - name: Validate wheels + run: twine check dist-publish/*.whl + + - name: Publish to TestPyPI + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.TESTPYPI_API_TOKEN }} + run: | + twine upload --repository testpypi --skip-existing dist-publish/*.whl + + nightly-test: + needs: version-stamp + runs-on: ubuntu-22.04 + env: + CI: "true" + SCCACHE_GHA_ENABLED: "true" + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + + - name: Set up Python 3.12 + uses: actions/setup-python@v5 + with: + python-version: '3.12' + + - name: Install Rust toolchain + uses: dtolnay/rust-toolchain@stable + + - name: Install and start etcd + run: | + wget -q https://github.com/etcd-io/etcd/releases/download/v3.6.1/etcd-v3.6.1-linux-amd64.tar.gz + tar xzf etcd-v3.6.1-linux-amd64.tar.gz + sudo mv etcd-v3.6.1-linux-amd64/etcd* /usr/local/bin/ + etcd --advertise-client-urls http://127.0.0.1:2379 --listen-client-urls http://127.0.0.1:2379 & + sleep 3 + ETCDCTL_API=3 etcdctl --endpoints=http://127.0.0.1:2379 endpoint health + + - name: Free up disk space + run: | + sudo rm -rf /usr/share/dotnet /opt/ghc /opt/hostedtoolcache/CodeQL /usr/local/lib/android + df -h + + - name: Install CUDA Toolkit + uses: Jimver/cuda-toolkit@v0.2.24 + with: + cuda: '12.8.1' + linux-local-args: '["--toolkit"]' + method: 'network' + sub-packages: '["nvcc"]' + + - name: Install build utilities + run: | + sudo apt-get update + sudo apt-get install -y ninja-build + + - name: Run sccache-cache + uses: mozilla-actions/sccache-action@v0.0.9 + + - name: Configure sccache + uses: actions/github-script@v7 + with: + script: | + core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); + core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); + + - name: Configure project (Release, no ASAN) + run: | + sudo apt update -y + sudo bash -x dependencies.sh -y + echo "/usr/local/go/bin" >> "$GITHUB_PATH" + mkdir build && cd build + cmake -G Ninja .. \ + -DUSE_HTTP=ON -DUSE_CXL=ON -DUSE_UB=ON -DUSE_ETCD=ON -DUSE_CUDA=ON \ + -DWITH_P2P_STORE=ON \ + -DSTORE_USE_ETCD=ON -DCMAKE_BUILD_TYPE=Release \ + -DBUILD_UNIT_TESTS=ON -DENABLE_SCCACHE=ON + + - name: Configure CUDA driver runtime + run: | + cuda_driver_library=$(sed -n \ + 's/^CUDA_cuda_driver_LIBRARY:FILEPATH=//p' build/CMakeCache.txt) + if [ -z "$cuda_driver_library" ] || [ ! -f "$cuda_driver_library" ]; then + echo "::error::CMake did not resolve the CUDA driver library" + exit 1 + fi + + cuda_driver_dir=$(dirname "$cuda_driver_library") + if [ ! -e "$cuda_driver_dir/libcuda.so.1" ]; then + sudo ln -s "$(basename "$cuda_driver_library")" \ + "$cuda_driver_dir/libcuda.so.1" + fi + echo "LIBRARY_PATH=$cuda_driver_dir:${LIBRARY_PATH:-}" >> "$GITHUB_ENV" + echo "LD_LIBRARY_PATH=$cuda_driver_dir:${LD_LIBRARY_PATH:-}" >> "$GITHUB_ENV" + + - name: Build project + run: | + cd build + cmake --build . -j$(nproc) + sudo -E cmake --install . + + - name: Build nvlink_allocator.so + run: | + mkdir -p build/mooncake-transfer-engine/nvlink-allocator + cd mooncake-transfer-engine/nvlink-allocator + bash build.sh ../../build/mooncake-transfer-engine/nvlink-allocator/ + + - name: Start Metadata Server + run: | + cd mooncake-transfer-engine/example/http-metadata-server-python + pip install aiohttp + python ./bootstrap_server.py & + echo "NIGHTLY_METADATA_SERVER_PID=$!" >> "$GITHUB_ENV" + sleep 2 + + - name: Run CTest unit tests + run: | + # Keep ephemeral TIME_WAIT sockets off mooncake_client's RPC port. + reserved_ports=$(sysctl -n net.ipv4.ip_local_reserved_ports) + sudo sysctl -w "net.ipv4.ip_local_reserved_ports=${reserved_ports:+$reserved_ports,}50052" + 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 --parallel $(nproc) --output-on-failure + + - name: Run Mooncake Store Rust smoke test + env: + MOONCAKE_STORE_CLUSTER_ID: nightly_rust_cluster + MOONCAKE_STORE_RUST_LINK_ASAN: "0" + run: ./scripts/ci/run_store_rust_smoke.sh + + - name: Run Transfer Engine Rust smoke test + env: + MOONCAKE_TE_RUST_LINK_ASAN: "0" + run: ./scripts/ci/run_transfer_engine_rust_smoke.sh + + - name: Check Mooncake Store Rust bindings, examples, and tests + env: + MOONCAKE_BUILD_DIR: ${{ github.workspace }}/build + MOONCAKE_STORE_LIB_DIR: ${{ github.workspace }}/build/mooncake-store/src + MOONCAKE_STORE_INCLUDE_DIR: ${{ github.workspace }}/mooncake-store/include + run: | + export LD_LIBRARY_PATH="$GITHUB_WORKSPACE/build/mooncake-asio:$GITHUB_WORKSPACE/build/mooncake-store/src:$GITHUB_WORKSPACE/build/mooncake-store/src/cachelib_memory_allocator:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base:$GITHUB_WORKSPACE/build/mooncake-common/etcd:/usr/local/lib:${LD_LIBRARY_PATH:-}" + cd mooncake-store/rust + cargo test --lib + cargo test --examples --tests --no-run + cargo clean + + - name: Verify Mooncake Store Rust dlopen bindings and packaging + env: + MOONCAKE_STORE_INCLUDE_DIR: ${{ github.workspace }}/mooncake-store/include + run: | + cd mooncake-store/rust + cargo run --locked --example generate_dlopen_bindings + git diff --exit-code -- src/generated/ffi_dlopen_bindings.rs + cargo package --no-verify --allow-dirty + crate=$(ls target/package/mooncake_store-*.crate | head -1) + dest=$(mktemp -d) + tar xzf "$crate" -C "$dest" + (cd "$dest"/mooncake_store-* && cargo check --no-default-features --features dlopen) + + - name: Run Go store binding integration tests + env: + MOONCAKE_STORE_CLUSTER_ID: nightly_go_cluster + MOONCAKE_STORE_GO_SANITIZED: "0" + run: ./scripts/ci/run_store_go_integration.sh + + - name: Build and install Python wheel for integration tests + run: | + python -m venv "${RUNNER_TEMP}/mooncake-nightly-venv" + echo "${RUNNER_TEMP}/mooncake-nightly-venv/bin" >> "$GITHUB_PATH" + source "${RUNNER_TEMP}/mooncake-nightly-venv/bin/activate" + export LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:/usr/local/lib + export CUDA_HOME=/usr/local/cuda + PYTHON_VERSION=3.12 OUTPUT_DIR=dist ./scripts/build_wheel.sh + python -m pip install mooncake-wheel/dist/*.whl + + - name: Run Python integration tests (full suite) + env: + MC_METADATA_SERVER: http://127.0.0.1:8080/metadata + RUN_TESTS_METADATA_SERVER_MODE: external + DEFAULT_KV_LEASE_TTL: "500" + TEST_CXL: "1" + run: | + export LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:/usr/local/lib + bash scripts/run_tests.sh + + - name: Run Python API end-to-end tests + 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: | + export LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:/usr/local/lib + kill "$NIGHTLY_METADATA_SERVER_PID" 2>/dev/null || true + for _ in {1..50}; do + if ! kill -0 "$NIGHTLY_METADATA_SERVER_PID" 2>/dev/null; then + break + fi + sleep 0.1 + done + if kill -0 "$NIGHTLY_METADATA_SERVER_PID" 2>/dev/null; then + kill -KILL "$NIGHTLY_METADATA_SERVER_PID" + fi + mkdir -p /tmp/mooncake_storage + mooncake_master \ + --default_kv_lease_ttl=500 \ + --eviction_high_watermark_ratio=0.95 \ + --cluster_id=nightly_test_cluster \ + --port 50051 \ + --enable_http_metadata_server=true & + master_pid=$! + trap 'kill "$master_pid" 2>/dev/null || true; wait "$master_pid" 2>/dev/null || true' EXIT + sleep 3 + python scripts/test_tensor_api.py -n 1 + python scripts/test_async_store.py + python scripts/test_copy_move_api.py + python scripts/test_drain_http_api.py --timeout-sec 90 + + - name: Run RPC Communicator Bandwidth Test + run: | + rpc_server_pid="" + cleanup() { + status=$? + trap - EXIT + if [ -n "$rpc_server_pid" ]; then + kill "$rpc_server_pid" 2>/dev/null || true + wait "$rpc_server_pid" 2>/dev/null || true + fi + if [ "$status" -ne 0 ]; then + cat "$RUNNER_TEMP/rpc-server.log" 2>/dev/null || true + fi + exit "$status" + } + trap cleanup EXIT + + python -u mooncake-transfer-engine/tests/rpc_communicator_test.py \ + server --url 127.0.0.1:9004 --data-size 1 \ + >"$RUNNER_TEMP/rpc-server.log" 2>&1 & + rpc_server_pid=$! + + rpc_ready=false + for _ in {1..50}; do + if ! kill -0 "$rpc_server_pid" 2>/dev/null; then + cat "$RUNNER_TEMP/rpc-server.log" + echo "::error::RPC communicator server exited before becoming ready" + exit 1 + fi + if ss -H -ltn 'sport = :9004' | grep -q .; then + rpc_ready=true + break + fi + sleep 0.1 + done + if [ "$rpc_ready" != true ]; then + cat "$RUNNER_TEMP/rpc-server.log" + echo "::error::RPC communicator server did not become ready within 5 seconds" + exit 1 + fi + + client_rc=0 + timeout 10 python -u \ + mooncake-transfer-engine/tests/rpc_communicator_test.py \ + client --url 127.0.0.1:9004 --threads 2 --data-size 1 \ + >"$RUNNER_TEMP/rpc-client.log" 2>&1 || \ + client_rc=$? + cat "$RUNNER_TEMP/rpc-client.log" + if [ "$client_rc" -ne 0 ] && [ "$client_rc" -ne 124 ]; then + echo "::error::RPC communicator client failed with exit code $client_rc" + exit "$client_rc" + fi + if ! grep -q '^bandwidth:' "$RUNNER_TEMP/rpc-client.log"; then + cat "$RUNNER_TEMP/rpc-server.log" + echo "::error::RPC communicator did not complete a successful transfer" + exit 1 + fi + shell: bash + + nightly-coverage: + runs-on: ubuntu-22.04 + env: + CI: "true" + SCCACHE_GHA_ENABLED: "true" + CXXFLAGS: --coverage + CFLAGS: --coverage + LDFLAGS: --coverage + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + + - name: Set up Python 3.12 + uses: actions/setup-python@v5 + with: + python-version: '3.12' + + - name: Install Rust toolchain + uses: dtolnay/rust-toolchain@stable + + - name: Install and start etcd + run: | + wget -q https://github.com/etcd-io/etcd/releases/download/v3.6.1/etcd-v3.6.1-linux-amd64.tar.gz + tar xzf etcd-v3.6.1-linux-amd64.tar.gz + sudo mv etcd-v3.6.1-linux-amd64/etcd* /usr/local/bin/ + etcd --advertise-client-urls http://127.0.0.1:2379 --listen-client-urls http://127.0.0.1:2379 & + sleep 3 + ETCDCTL_API=3 etcdctl --endpoints=http://127.0.0.1:2379 endpoint health + + - name: Free up disk space + uses: ./.github/actions/free-disk-space + + - name: Install CUDA Toolkit + uses: Jimver/cuda-toolkit@v0.2.24 + with: + cuda: '12.8.1' + linux-local-args: '["--toolkit"]' + method: 'network' + sub-packages: '["nvcc"]' + + - name: Install coverage tools and build utilities + run: | + sudo apt-get update + sudo apt-get install -y lcov gcovr ninja-build + + - name: Test HugeTLB sizing helper + run: python3 scripts/test_hicache_hugepage_requirements.py + + - name: Run sccache-cache + uses: mozilla-actions/sccache-action@v0.0.9 + + - name: Configure sccache + uses: actions/github-script@v7 + with: + script: | + core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); + core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); + + - name: Configure project with coverage support + run: | + sudo apt update -y + 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 -DENABLE_DEBUG_SYMBOLS=OFF + + - name: Build project + run: | + cd build + cmake --build . --parallel "$(nproc)" + sudo cmake --install . + + - name: Start Metadata Server + run: | + cd mooncake-transfer-engine/example/http-metadata-server-python + pip install aiohttp + python ./bootstrap_server.py & + sleep 2 + + - name: Run CTest with coverage + run: | + 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 --parallel $(nproc) --output-on-failure + + - name: Run Mooncake Store Rust sanitizer smoke test + env: + MOONCAKE_STORE_CLUSTER_ID: nightly_asan_rust_cluster + MOONCAKE_STORE_RUST_LINK_ASAN: "1" + run: ./scripts/ci/run_store_rust_smoke.sh + + - name: Run Transfer Engine Rust sanitizer smoke test + env: + MOONCAKE_TE_RUST_LINK_ASAN: "1" + run: ./scripts/ci/run_transfer_engine_rust_smoke.sh + + - name: Run Go store binding sanitizer integration tests + env: + MOONCAKE_STORE_CLUSTER_ID: nightly_asan_go_cluster + MOONCAKE_STORE_GO_SANITIZED: "1" + run: ./scripts/ci/run_store_go_integration.sh + + - name: Generate coverage report + id: coverage + run: | + cd build + echo "=== Starting coverage report generation ===" + echo "Current directory: $(pwd)" + + echo "=== Looking for .gcda files ===" + find . -name "*.gcda" 2>/dev/null | head -10 || echo "No .gcda files found" + + echo "=== Running lcov ===" + lcov --capture --directory . --output-file coverage.info 2>&1 || { + echo "WARNING: lcov failed to capture coverage data" + echo "Creating minimal lcov-compliant coverage file to allow CI to continue" + echo "TN:dummy" > coverage.filtered.info + echo "SF:/dev/null" >> coverage.filtered.info + echo "DA:0,0" >> coverage.filtered.info + echo "end_of_record" >> coverage.filtered.info + echo "coverage_failed=true" >> "$GITHUB_OUTPUT" + exit 0 + } + + echo "=== Processing coverage data ===" + lcov --remove coverage.info '/usr/*' '*/test/*' '*/third_party/*' '*/benchmarks/*' --output-file coverage.filtered.info 2>&1 || true + + echo "=== Generating HTML report ===" + genhtml coverage.filtered.info --output-directory coverage_report 2>&1 || echo "genhtml failed, continuing..." + + echo "=== Coverage summary ===" + lcov --list coverage.filtered.info 2>&1 || echo "lcov list failed" + + echo "=== Coverage report generation completed ===" + + - name: Upload coverage to Codecov + uses: codecov/codecov-action@v4 + with: + files: build/coverage.filtered.info + flags: unittests + name: nightly-code-coverage-report + token: ${{ secrets.CODECOV_TOKEN }} + fail_ci_if_error: false + continue-on-error: true + + - name: Check coverage status + if: always() + run: | + if [ "${{ steps.coverage.outputs.coverage_failed }}" = "true" ]; then + echo "Coverage collection failed but the nightly workflow continued" + echo "::warning::Code coverage collection failed. Please check the logs." + else + echo "Coverage collected successfully" + fi + + nightly-gate: + # always() bypasses skipped dependencies, so this job needs its own guard. + if: ${{ always() && github.repository == 'kvcache-ai/Mooncake' }} + needs: + - version-stamp + - build-wheels + - tone-sglang-integration + - build-musa + - build-docker + - ascend-test + - publish-testpypi + - nightly-test + - nightly-coverage + runs-on: ubuntu-22.04 + steps: + - name: Check results + run: | + echo "=== Nightly Gate Summary ===" + echo "$NEEDS_JSON" | jq -r ' + to_entries[] | "\(.key): \(.value.result)"' + + failing=$(echo "$NEEDS_JSON" | jq -r \ + --arg allow_publish_skip "$ALLOW_PUBLISH_SKIP" ' + to_entries[] | + select( + .value.result != "success" and + (.key != "publish-testpypi" or + .value.result != "skipped" or + $allow_publish_skip != "true") + ) | + "\(.key): \(.value.result)"') + if [ -n "$failing" ]; then + echo "::error::The following nightly jobs did not succeed:" + echo "$failing" + exit 1 + fi + echo "All nightly jobs passed!" + env: + NEEDS_JSON: ${{ toJSON(needs) }} + ALLOW_PUBLISH_SKIP: ${{ github.event_name == 'workflow_dispatch' && github.event.inputs.skip_publish == 'true' }} + + notify-failure: + if: ${{ always() && needs.nightly-gate.result == 'failure' }} + needs: [nightly-gate] + runs-on: ubuntu-22.04 + permissions: + issues: write + steps: + - name: Create failure issue + uses: actions/github-script@v7 + with: + script: | + const labelName = 'nightly-failure'; + // Ensure the label exists + try { + await github.rest.issues.getLabel({ + owner: context.repo.owner, + repo: context.repo.repo, + name: labelName, + }); + } catch (e) { + if (e.status === 404) { + await github.rest.issues.createLabel({ + owner: context.repo.owner, + repo: context.repo.repo, + name: labelName, + color: 'd73a4a', + description: 'Nightly CI workflow failure', + }); + } + } + const title = `Nightly build/test failure - ${new Date().toISOString().slice(0, 10)}`; + const body = [ + '## Nightly Failure Report', + '', + `**Run**: ${context.serverUrl}/${context.repo.owner}/${context.repo.repo}/actions/runs/${context.runId}`, + `**Branch**: ${context.ref}`, + `**Timestamp**: ${new Date().toISOString()}`, + '', + 'Please investigate the failed jobs in the workflow run linked above.', + ].join('\n'); + await github.rest.issues.create({ + owner: context.repo.owner, + repo: context.repo.repo, + title, + body, + labels: [labelName], + }); diff --git a/.github/workflows/pr-tidy.yml b/.github/workflows/pr-tidy.yml new file mode 100644 index 0000000000..8234296304 --- /dev/null +++ b/.github/workflows/pr-tidy.yml @@ -0,0 +1,31 @@ +name: PR Tidy + +on: + pull_request_target: + types: [opened] + +permissions: + pull-requests: write + +jobs: + tidy: + runs-on: ubuntu-latest + steps: + - name: Clean PR description + uses: actions/github-script@v7 + with: + script: | + const pr = context.payload.pull_request; + if (!pr.body) return; + + // Remove HTML comment blocks (template instructions) + let body = pr.body.replace(//g, '').trim(); + + if (body !== pr.body) { + await github.rest.pulls.update({ + owner: context.repo.owner, + repo: context.repo.repo, + pull_number: pr.number, + body: body + }); + } diff --git a/.github/workflows/pre-release.yaml b/.github/workflows/pre-release.yaml new file mode 100644 index 0000000000..85f6f96199 --- /dev/null +++ b/.github/workflows/pre-release.yaml @@ -0,0 +1,269 @@ +name: Pre-Release + +# Build the same core wheel matrix as the stable release workflows, publish the +# exact artifacts to TestPyPI, then consume them in smoke and T-one tests. +# No production PyPI or GitHub Release publication happens here. +on: + push: + tags: + - 'v*-rc*' + - 'v*-alpha*' + - 'v*-beta*' + - 'v*-pre*' + +jobs: + version-stamp: + name: Normalize pre-release version + if: github.repository == 'kvcache-ai/Mooncake' + runs-on: ubuntu-22.04 + permissions: + contents: read + outputs: + package_version: ${{ steps.version.outputs.package_version }} + steps: + - name: Checkout release gate + uses: actions/checkout@v4 + with: + persist-credentials: false + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.12' + + - name: Install version parser + run: python -m pip install packaging + + - name: Normalize tag as a PEP 440 pre-release + id: version + env: + RELEASE_TAG: ${{ github.ref_name }} + run: | + normalized=$(python scripts/ci/testpypi_wheel_gate.py \ + normalize-version --tag "$RELEASE_TAG") + echo "package_version=${normalized}" >> "$GITHUB_OUTPUT" + echo "Normalized ${RELEASE_TAG} to ${normalized}" + + build: + needs: version-stamp + strategy: + fail-fast: false + matrix: + include: + - variant: cuda + architecture: x86_64 + artifact-prefix: mooncake-wheel-pre-release + - variant: cuda + architecture: arm64 + artifact-prefix: mooncake-wheel-arm64-pre-release + - variant: cuda13 + architecture: x86_64 + artifact-prefix: mooncake-wheel-cuda13-pre-release + - variant: cuda13 + architecture: arm64 + artifact-prefix: mooncake-wheel-cuda13-arm64-pre-release + - variant: non-cuda + architecture: x86_64 + artifact-prefix: mooncake-wheel-non-cuda-pre-release + - variant: non-cuda + architecture: arm64 + artifact-prefix: mooncake-wheel-non-cuda-arm64-pre-release + uses: ./.github/workflows/_build-wheel.yaml + with: + variant: ${{ matrix.variant }} + architecture: ${{ matrix.architecture }} + artifact-prefix: ${{ matrix.artifact-prefix }} + version-override: ${{ needs.version-stamp.outputs.package_version }} + + publish-testpypi: + name: Validate and publish wheels to TestPyPI + needs: [version-stamp, build] + runs-on: ubuntu-22.04 + environment: nightly + concurrency: + group: pre-release-testpypi-${{ needs.version-stamp.outputs.package_version }} + cancel-in-progress: false + permissions: + contents: read + env: + PACKAGE_VERSION: ${{ needs.version-stamp.outputs.package_version }} + steps: + - name: Checkout source + uses: actions/checkout@v4 + + - name: Download all pre-release wheel artifacts + uses: actions/download-artifact@v4 + with: + path: mooncake-wheel/dist-all + pattern: mooncake-wheel*pre-release* + + - name: Collect wheels for publication + run: | + mkdir -p mooncake-wheel/dist-release + find mooncake-wheel/dist-all -name "*.whl" \ + -exec cp {} mooncake-wheel/dist-release/ \; + echo "Pre-release tag: ${GITHUB_REF_NAME}" + echo "Package version: ${PACKAGE_VERSION}" + echo "Collected wheels:" + ls -la mooncake-wheel/dist-release/ + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.12' + + - name: Install publication tools + run: python -m pip install packaging twine + + - name: Require TestPyPI credentials + env: + TESTPYPI_API_TOKEN: ${{ secrets.TESTPYPI_API_TOKEN }} + run: | + if [ -z "$TESTPYPI_API_TOKEN" ]; then + echo "TESTPYPI_API_TOKEN is required" + exit 1 + fi + + - name: Validate wheels with twine + run: twine check mooncake-wheel/dist-release/*.whl + + - name: Validate empty or resumable TestPyPI upload state + run: | + python scripts/ci/testpypi_wheel_gate.py validate-upload-state \ + --directory mooncake-wheel/dist-release \ + --version "$PACKAGE_VERSION" + + - name: Publish exact wheels to TestPyPI + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.TESTPYPI_API_TOKEN }} + run: >- + twine upload --repository testpypi --skip-existing + mooncake-wheel/dist-release/*.whl + + - name: Verify every published wheel hash + run: | + python scripts/ci/testpypi_wheel_gate.py wait-upload-state \ + --directory mooncake-wheel/dist-release \ + --version "$PACKAGE_VERSION" + + consume-testpypi: + name: Consume ${{ matrix.package }} on ${{ matrix.architecture }} + needs: [version-stamp, publish-testpypi] + runs-on: ${{ matrix.architecture == 'aarch64' && 'ubuntu-22.04-arm' || 'ubuntu-22.04' }} + permissions: + contents: read + strategy: + fail-fast: false + matrix: + package: + - mooncake-transfer-engine + - mooncake-transfer-engine-cuda13 + - mooncake-transfer-engine-non-cuda + architecture: + - x86_64 + - aarch64 + env: + PACKAGE_VERSION: ${{ needs.version-stamp.outputs.package_version }} + TARGET_PACKAGE: ${{ matrix.package }} + steps: + - name: Install CUDA runtime and driver stub + if: ${{ matrix.package != 'mooncake-transfer-engine-non-cuda' }} + env: + CUDA_PACKAGE_VERSION: ${{ matrix.package == 'mooncake-transfer-engine-cuda13' && '13-0' || '12-8' }} + run: | + set -euo pipefail + if [ "$(uname -m)" = x86_64 ]; then + repository_architecture=x86_64 + else + repository_architecture=sbsa + fi + keyring=$(mktemp --suffix=.deb) + curl --fail --location --retry 3 \ + "https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/${repository_architecture}/cuda-keyring_1.1-1_all.deb" \ + --output "$keyring" + sudo dpkg -i "$keyring" + sudo apt-get update + sudo apt-get install -y \ + "cuda-cudart-${CUDA_PACKAGE_VERSION}" \ + "cuda-driver-dev-${CUDA_PACKAGE_VERSION}" + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.12' + + - name: Download, install, and smoke the indexed wheel + run: | + set -euo pipefail + sudo apt-get update + sudo apt-get install -y \ + libcurl4 libibverbs1 rdma-core librdmacm1 libnuma1 liburing2 + + download_dir=$(mktemp -d) + + python -m pip download \ + --index-url https://test.pypi.org/simple \ + --no-cache-dir \ + --no-deps \ + --only-binary=:all: \ + --pre \ + --dest "$download_dir" \ + "${TARGET_PACKAGE}==${PACKAGE_VERSION}" + + # Install the wheel downloaded above. Its ordinary dependencies are + # resolved from production PyPI, never TestPyPI. + python -m pip install \ + --index-url https://pypi.org/simple \ + "$download_dir"/*.whl + + if [ "$TARGET_PACKAGE" != mooncake-transfer-engine-non-cuda ]; then + cuda_stub= + for candidate in /usr/local/cuda*/lib64/stubs/libcuda.so \ + /usr/local/cuda*/targets/*/lib/stubs/libcuda.so; do + if [ -f "$candidate" ]; then + cuda_stub="$candidate" + break + fi + done + if [ -z "$cuda_stub" ]; then + echo "CUDA driver stub libcuda.so was not found" + exit 1 + fi + stub_dir=$(mktemp -d) + ln -s "$cuda_stub" "$stub_dir/libcuda.so.1" + export LD_LIBRARY_PATH="$stub_dir:${LD_LIBRARY_PATH:-}" + for directory in /usr/local/cuda*/lib64 \ + /usr/local/cuda*/targets/*/lib; do + if [ -d "$directory" ]; then + export LD_LIBRARY_PATH="$directory:$LD_LIBRARY_PATH" + fi + done + fi + + python - "$TARGET_PACKAGE" "$PACKAGE_VERSION" <<'PY' + import importlib.metadata + import sys + + package, expected = sys.argv[1:] + installed = importlib.metadata.version(package) + if installed != expected: + raise SystemExit( + f"installed {package} version {installed}, expected {expected}" + ) + + import mooncake + import mooncake.engine + import mooncake.http_metadata_server + import mooncake.store + PY + mooncake_http_metadata_server --help >/dev/null + mooncake_master --version + + tone-integration: + name: Validate T-one integration with TestPyPI wheel + needs: [version-stamp, consume-testpypi] + uses: ./.github/workflows/integration-test.yml + with: + testpypi_version: ${{ needs.version-stamp.outputs.package_version }} + secrets: inherit diff --git a/.github/workflows/publish-master-image.yaml b/.github/workflows/publish-master-image.yaml new file mode 100644 index 0000000000..9a62b6e6b1 --- /dev/null +++ b/.github/workflows/publish-master-image.yaml @@ -0,0 +1,255 @@ +name: Publish Master Image + +# Manual trigger only: this installs an already-published PyPI wheel, so it must be +# run after that version's wheel is published — avoids racing release.yaml, which builds +# and publishes the wheel in the same v* tag event. +# +# Two CUDA flavors are published from the same version input, because the wheel ships as +# two PyPI projects that are versioned in lockstep: +# cuda12 -> mooncake-transfer-engine -> : (and :latest) +# cuda13 -> mooncake-transfer-engine-cuda13 -> :-cuda13 (and :latest-cuda13) +# The cuda12 tag stays unsuffixed so existing pulls keep working. Each flavor is built +# from its own Dockerfile (docker/master.Dockerfile, docker/master-cuda13.Dockerfile); +# those two files must differ only in the CUDA-flavor lines, which `prepare` enforces. +# +# Publish flow: validate input -> check the Dockerfiles are in sync -> per flavor: confirm +# amd64+arm64 wheels exist -> build multi-arch and push the :[-cuda13] tag -> +# smoke-test amd64 AND arm64 -> only then promote :latest[-cuda13] (when requested). +# :latest is copied from the smoked : digest, so the default pull tag can never +# point at an image that failed smoke. A failed smoke still leaves the : tag +# public; since this is manual, delete it from Docker Hub by hand. The flavors run as +# independent matrix legs with fail-fast disabled, so one flavor failing never cancels +# the other mid-push. +on: + workflow_dispatch: + inputs: + mooncake_version: + description: "Published wheel version (also the image tag), e.g. 0.3.11.post1. Both PyPI projects use the same version." + required: true + type: string + flavors: + description: "Which CUDA flavor(s) to publish. Keep 'both' unless backfilling a single flavor." + required: false + default: both + type: choice + options: + - both + - cuda12 + - cuda13 + tag_latest: + description: "Also move :latest (per flavor) to this build. Only enable when publishing the newest release (a backfill/retry of an older wheel must NOT move :latest)." + required: false + default: false + type: boolean + +permissions: + contents: read + +# Serialize publishes so two concurrent runs can't interleave and move :latest backwards. +concurrency: + group: publish-master-image + cancel-in-progress: false + +jobs: + # Cheap gate: everything that is flavor-independent runs once, before any runner spends + # time on QEMU/buildx. Also emits the flavor matrix consumed by publish-master. + prepare: + runs-on: ubuntu-22.04 + outputs: + matrix: ${{ steps.matrix.outputs.matrix }} + env: + # inputs.* placed in env (not interpolated into shell source) — injection-safe. + MOONCAKE_VERSION: ${{ inputs.mooncake_version }} + FLAVORS: ${{ inputs.flavors }} + steps: + - name: Checkout + uses: actions/checkout@v4 + with: + persist-credentials: false + + - name: Validate mooncake_version + run: | + if ! printf '%s' "$MOONCAKE_VERSION" | grep -Eq '^[0-9]+\.[0-9]+\.[0-9]+(\.post[0-9]+)?$'; then + echo "Invalid mooncake_version '$MOONCAKE_VERSION': expected e.g. 0.3.11 or 0.3.11.post1" >&2 + exit 1 + fi + + # The two flavors are the same image with a different CUDA runtime, so the two + # Dockerfiles are kept byte-identical apart from the CUDA-flavor lines. Normalize + # those lines and diff: any other divergence (a runtime package added to one file + # only, say) fails here instead of silently shipping one broken flavor. + - name: Check the two master Dockerfiles are in sync + run: | + norm() { + sed -E \ + -e 's#^FROM nvidia/cuda:.* AS cudalibs$#FROM AS cudalibs#' \ + -e 's#libcudart\.so\.[0-9]+#libcudart.so.#g' \ + -e 's#^ mooncake-transfer-engine(-cuda13)?==# ==#' \ + "$1" + } + if ! diff -u <(norm docker/master.Dockerfile) <(norm docker/master-cuda13.Dockerfile); then + echo "::error::docker/master.Dockerfile and docker/master-cuda13.Dockerfile differ outside the CUDA-flavor lines" >&2 + exit 1 + fi + echo "Dockerfiles in sync (only the CUDA-flavor lines differ)" + + - name: Build flavor matrix + id: matrix + run: | + python3 - <<'PY' >> "$GITHUB_OUTPUT" + import json, os, sys + FLAVORS = { + "cuda12": { + "flavor": "cuda12", + "package": "mooncake-transfer-engine", + "dockerfile": "docker/master.Dockerfile", + "tag_suffix": "", + }, + "cuda13": { + "flavor": "cuda13", + "package": "mooncake-transfer-engine-cuda13", + "dockerfile": "docker/master-cuda13.Dockerfile", + "tag_suffix": "-cuda13", + }, + } + sel = os.environ["FLAVORS"] + keys = list(FLAVORS) if sel == "both" else [sel] + if any(k not in FLAVORS for k in keys): + sys.exit(f"Unknown flavors input: {sel!r}") + print("matrix=" + json.dumps([FLAVORS[k] for k in keys])) + PY + cat "$GITHUB_OUTPUT" + + publish-master: + needs: prepare + runs-on: ubuntu-22.04 + strategy: + # One flavor failing must not cancel the other: a cancelled leg could be killed + # between its push and its smoke test, leaving an untested tag with no verdict. + fail-fast: false + matrix: + include: ${{ fromJSON(needs.prepare.outputs.matrix) }} + env: + REPO: docker.io/kvcacheai/mooncake + # inputs.*/matrix.* placed in env (not interpolated into shell source) — injection-safe. + MOONCAKE_VERSION: ${{ inputs.mooncake_version }} + MOONCAKE_PACKAGE: ${{ matrix.package }} + IMAGE_TAG: ${{ inputs.mooncake_version }}${{ matrix.tag_suffix }} + TAG_SUFFIX: ${{ matrix.tag_suffix }} + # Pinned for the official image: provenance must always be the public PyPI index. + # Debug builds against another index belong in a separate workflow that does NOT + # push official kvcacheai/mooncake tags. + PIP_INDEX_URL: https://pypi.org/simple + steps: + - name: Checkout + uses: actions/checkout@v4 + with: + persist-credentials: false + + # This image is intentionally the CPython 3.12 flavor: both master Dockerfiles pin + # python:3.12-slim-trixie, so only cp312 wheels are ABI-compatible. If that base + # Python is ever bumped (e.g. to 3.13), update the "cp312" match below to match — + # otherwise this preflight would green-light wheels the image cannot import. + - name: Confirm cp312 amd64+arm64 wheels are published on PyPI + run: | + python3 - <<'PY' + import json, os, sys, urllib.request + v = os.environ["MOONCAKE_VERSION"] + pkg = os.environ["MOONCAKE_PACKAGE"] + url = f"https://pypi.org/pypi/{pkg}/{v}/json" + try: + data = json.load(urllib.request.urlopen(url, timeout=30)) + except Exception as e: + sys.exit(f"Cannot fetch PyPI metadata for {pkg} {v}: {e}") + arch_ok = {"x86_64": False, "aarch64": False} + for f in data.get("urls", []): + fn = f.get("filename", "") + if fn.endswith(".whl") and "cp312" in fn: + for arch in arch_ok: + if arch in fn: + arch_ok[arch] = True + missing = [a for a, ok in arch_ok.items() if not ok] + if missing: + sys.exit(f"Missing cp312 wheel(s) for {pkg} {v}: {missing}. " + "Both amd64 (x86_64) and arm64 (aarch64) must be published first.") + print(f"cp312 amd64+arm64 wheels present for {pkg} {v}") + PY + + - name: Free up disk space + uses: ./.github/actions/free-disk-space + + - name: Set up QEMU + uses: docker/setup-qemu-action@v3 + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Log in to Docker Hub + uses: docker/login-action@v3 + with: + registry: docker.io + # For an org access token (OAT), the login username is the org name + # (Docker: docker login --username ) — not a secret. + username: kvcacheai + password: ${{ secrets.DOCKERHUB_TOKEN }} + + # Build multi-arch and push ONLY the immutable :[-cuda13] tag. :latest is + # promoted in a later step, after smoke passes. provenance/sbom disabled to keep the + # index free of unknown/unknown attestation entries (they would otherwise show up in + # `imagetools inspect`). + - name: Build and push :${{ matrix.tag_suffix }} (multi-arch) + id: build + uses: docker/build-push-action@v6 + with: + context: . + file: ${{ matrix.dockerfile }} + platforms: linux/amd64,linux/arm64 + provenance: false + sbom: false + build-args: | + MOONCAKE_VERSION=${{ inputs.mooncake_version }} + PIP_INDEX_URL=${{ env.PIP_INDEX_URL }} + tags: ${{ env.REPO }}:${{ inputs.mooncake_version }}${{ matrix.tag_suffix }} + push: true + + # Smoke tests run AFTER the : push (a failed smoke leaves : public + # — delete it from Docker Hub by hand; this manual workflow wires no auto-cleanup). + # :latest is NOT pushed yet, so a bad build can never poison the default pull tag. + # We exercise real entrypoints, not just imports: `mooncake_master --version` runs the + # compiled binary through the console-script/cli.py path and exits 0 (gflags + # SetVersionString) — note `--help` would exit 1 under gflags, so --version is used. + # `mooncake_http_metadata_server --help` (argparse) exits 0 and proves that entrypoint. + - name: Smoke test amd64 (entrypoints) + run: | + docker run --rm --platform linux/amd64 \ + "${REPO}:${IMAGE_TAG}" bash -c ' + set -euo pipefail + python3 -c "import mooncake.engine, mooncake.store" + mooncake_master --version + mooncake_http_metadata_server --help >/dev/null + echo "amd64 ok" + ' + + # arm64 is the riskier arch here: the aarch64 wheel is manylinux_2_39 (glibc >= 2.39), + # which is the whole reason the master Dockerfiles pin a trixie base. Smoke it under QEMU. + - name: Smoke test arm64 (entrypoints, QEMU) + run: | + docker run --rm --platform linux/arm64 \ + "${REPO}:${IMAGE_TAG}" bash -c ' + set -euo pipefail + python3 -c "import mooncake.engine, mooncake.store" + mooncake_master --version + mooncake_http_metadata_server --help >/dev/null + echo "arm64 ok" + ' + + # Promote :latest[-cuda13] ONLY after both arches pass smoke, and only when explicitly + # requested. imagetools create copies the tested : manifest by digest + # (no rebuild), so :latest can never point at an image that didn't pass smoke. + # Each flavor promotes its own latest tag: cuda12 -> :latest, cuda13 -> :latest-cuda13. + - name: Promote :latest${{ matrix.tag_suffix }} (post-smoke) + if: ${{ inputs.tag_latest }} + run: | + docker buildx imagetools create \ + --tag "${REPO}:latest${TAG_SUFFIX}" \ + "${REPO}:${IMAGE_TAG}" diff --git a/.github/workflows/release-cuda13.yaml b/.github/workflows/release-cuda13.yaml index 87c1b38ec4..4e45f477a3 100644 --- a/.github/workflows/release-cuda13.yaml +++ b/.github/workflows/release-cuda13.yaml @@ -5,141 +5,31 @@ on: tags: - 'v*' -env: - SCCACHE_GHA_ENABLED: "true" jobs: build: - runs-on: ubuntu-22.04 - permissions: - contents: write + if: ${{ !contains(github.ref_name, '-') }} strategy: + fail-fast: false matrix: - python-version: ['3.10', '3.11', '3.12', '3.13'] - env: - BUILD_WITH_EP: "1" - CU13_BUILD: "1" - TORCH_CUDA_ARCH_LIST: "8.0;9.0" - steps: - - name: Checkout source - uses: actions/checkout@v4 - - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - - - name: Free up disk space - run: | - sudo rm -rf /usr/share/dotnet - sudo rm -rf /opt/ghc - sudo rm -rf /opt/hostedtoolcache/CodeQL - sudo rm -rf /usr/local/lib/android - df -h - - - name: Install CUDA Toolkit 13 - uses: Jimver/cuda-toolkit@v0.2.29 - with: - cuda: '13.0.2' - linux-local-args: '["--toolkit"]' - method: 'network' - sub-packages: '["nvcc", "nvrtc-dev"]' - non-cuda-sub-packages: '["libcusparse-dev", "libcublas-dev", "libcusolver-dev"]' - - - name: Run sccache-cache - uses: mozilla-actions/sccache-action@v0.0.9 - - - name: Configure sccache - uses: actions/github-script@v7 - with: - script: | - core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); - core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); - - - name: Run sccache stat for check - shell: bash - run: ${SCCACHE_PATH} --show-stats - - - name: Configure project - run: | - sudo apt update -y - 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" -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release - shell: bash - - - name: Build project - run: | - export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH - cd build - make -j - sudo make install - shell: bash - - - name: Build nvlink_allocator.so - run: | - export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH - export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH - export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH - mkdir -p build/mooncake-transfer-engine/nvlink-allocator - cd mooncake-transfer-engine/nvlink-allocator - bash build.sh ../../build/mooncake-transfer-engine/nvlink-allocator/ - shell: bash - - - name: Generate Python version tag - id: generate_tag_release - run: | - echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT - shell: bash - - - name: Build Python wheel - run: | - # Set LD_LIBRARY_PATH for wheel building - export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib - PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh - env: - VERSION: ${{ env.VERSION }} - - - name: Upload Python wheel artifact - uses: actions/upload-artifact@v4 - with: - name: mooncake-wheel-cuda13-py${{ steps.generate_tag_release.outputs.python_version_tag }} - path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl + include: + - architecture: x86_64 + artifact-prefix: mooncake-wheel-cuda13 + - architecture: arm64 + artifact-prefix: mooncake-wheel-cuda13-arm64 + uses: ./.github/workflows/_build-wheel.yaml + with: + variant: cuda13 + architecture: ${{ matrix.architecture }} + artifact-prefix: ${{ matrix.artifact-prefix }} publish-release: + if: ${{ !contains(github.ref_name, '-') }} needs: build - runs-on: ubuntu-22.04 permissions: contents: write id-token: write - steps: - - name: Checkout source - uses: actions/checkout@v4 - - - name: Download all wheel artifacts - uses: actions/download-artifact@v4 - with: - path: mooncake-wheel/dist-all - pattern: mooncake-wheel-cuda13-py* - - - name: Prepare wheels for release - run: | - # Move all wheels to a single directory - mkdir -p mooncake-wheel/dist-release - find mooncake-wheel/dist-all -name "*.whl" -exec cp {} mooncake-wheel/dist-release/ \; - ls -la mooncake-wheel/dist-release/ - # List all collected wheels - echo "Collected wheels for release:" - ls -la mooncake-wheel/dist-release/ - - - name: Upload wheels to GitHub Release - uses: softprops/action-gh-release@v1 - with: - files: mooncake-wheel/dist-release/*.whl - - - name: Publish package to PyPI - if: github.repository == 'kvcache-ai/Mooncake' - uses: pypa/gh-action-pypi-publish@release/v1 - with: - packages-dir: mooncake-wheel/dist-release/ - password: ${{ secrets.PYPI_CU13_API_TOKEN }} + uses: ./.github/workflows/_publish-wheel.yaml + with: + artifact-pattern: 'mooncake-wheel-cuda13*' + secrets: + pypi-token: ${{ secrets.PYPI_CU13_API_TOKEN }} diff --git a/.github/workflows/release-efa-cuda13.yaml b/.github/workflows/release-efa-cuda13.yaml new file mode 100644 index 0000000000..1b0809b4d3 --- /dev/null +++ b/.github/workflows/release-efa-cuda13.yaml @@ -0,0 +1,36 @@ +name: Release EFA CUDA 13 + +on: + push: + tags: + - 'v*' + +# Publishes mooncake-transfer-engine-efa-cuda13, the CUDA 13-aware AWS EFA wheel. +jobs: + build: + if: ${{ !contains(github.ref_name, '-') }} + permissions: + contents: write + uses: ./.github/workflows/_build-efa-wheel.yaml + with: + variant: cuda13 + use-cuda: true + cuda-version: '13.0.2' + python-versions: '["3.10", "3.11", "3.12", "3.13"]' + build-profile: release + cmake-args: >- + -DUSE_HTTP=ON -DUSE_ETCD=ON -DWITH_EP=OFF -DSTORE_USE_ETCD=ON + variant-flag: EFA_CU13_BUILD + artifact-prefix: mooncake-wheel-efa-cuda13 + + publish-release: + if: ${{ !contains(github.ref_name, '-') }} + needs: build + permissions: + contents: write + id-token: write + uses: ./.github/workflows/_publish-wheel.yaml + with: + artifact-pattern: 'mooncake-wheel-efa-cuda13-py*' + secrets: + pypi-token: ${{ secrets.PYPI_API_TOKEN }} diff --git a/.github/workflows/release-efa-non-cuda.yaml b/.github/workflows/release-efa-non-cuda.yaml new file mode 100644 index 0000000000..6c6f8edc2b --- /dev/null +++ b/.github/workflows/release-efa-non-cuda.yaml @@ -0,0 +1,35 @@ +name: Release EFA Non-CUDA + +on: + push: + tags: + - 'v*' + +# Publishes mooncake-transfer-engine-efa-non-cuda, the CPU/DRAM AWS EFA wheel. +jobs: + build: + if: ${{ !contains(github.ref_name, '-') }} + permissions: + contents: write + uses: ./.github/workflows/_build-efa-wheel.yaml + with: + variant: non-cuda + use-cuda: false + python-versions: '["3.10", "3.11", "3.12", "3.13"]' + build-profile: release + cmake-args: >- + -DUSE_HTTP=ON -DUSE_ETCD=ON -DWITH_EP=OFF -DSTORE_USE_ETCD=ON + variant-flag: EFA_NON_CUDA_BUILD + artifact-prefix: mooncake-wheel-efa-non-cuda + + publish-release: + if: ${{ !contains(github.ref_name, '-') }} + needs: build + permissions: + contents: write + id-token: write + uses: ./.github/workflows/_publish-wheel.yaml + with: + artifact-pattern: 'mooncake-wheel-efa-non-cuda-py*' + secrets: + pypi-token: ${{ secrets.PYPI_API_TOKEN }} diff --git a/.github/workflows/release-efa.yaml b/.github/workflows/release-efa.yaml new file mode 100644 index 0000000000..a16aa5a7e8 --- /dev/null +++ b/.github/workflows/release-efa.yaml @@ -0,0 +1,36 @@ +name: Release EFA + +on: + push: + tags: + - 'v*' + +# Publishes mooncake-transfer-engine-efa, the CUDA 12-aware AWS EFA wheel. +jobs: + build: + if: ${{ !contains(github.ref_name, '-') }} + permissions: + contents: write + uses: ./.github/workflows/_build-efa-wheel.yaml + with: + variant: cuda + use-cuda: true + cuda-version: '12.8.1' + python-versions: '["3.10", "3.11", "3.12", "3.13"]' + build-profile: release + cmake-args: >- + -DUSE_HTTP=ON -DUSE_ETCD=ON -DWITH_EP=OFF -DSTORE_USE_ETCD=ON + variant-flag: EFA_BUILD + artifact-prefix: mooncake-wheel-efa + + publish-release: + if: ${{ !contains(github.ref_name, '-') }} + needs: build + permissions: + contents: write + id-token: write + uses: ./.github/workflows/_publish-wheel.yaml + with: + artifact-pattern: 'mooncake-wheel-efa-py*' + secrets: + pypi-token: ${{ secrets.PYPI_API_TOKEN }} diff --git a/.github/workflows/release-musa.yaml b/.github/workflows/release-musa.yaml new file mode 100644 index 0000000000..500e0e4aa2 --- /dev/null +++ b/.github/workflows/release-musa.yaml @@ -0,0 +1,212 @@ +name: Release MUSA + +on: + push: + tags: + - 'v*' + workflow_dispatch: + inputs: + tag: + description: 'Existing release tag to backfill (e.g. v0.3.12)' + required: true + type: string + +concurrency: + group: release-musa-${{ github.event_name == 'workflow_dispatch' && inputs.tag || github.ref_name }} + cancel-in-progress: false + +jobs: + build: + if: ${{ github.event_name == 'workflow_dispatch' || !contains(github.ref_name, '-') }} + runs-on: ubuntu-22.04 + container: registry.mthreads.com/mcconline/inference/pytorch:2.9.1.post1-py3.10-musa5.2.0-mp31-devel-ubuntu22.04-amd64 + + permissions: + contents: write + + strategy: + max-parallel: 2 + matrix: + python-version: ['3.10'] + + env: + MUSA_BUILD: "1" + + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ github.event_name == 'workflow_dispatch' && format('refs/tags/{0}', inputs.tag) || github.ref }} + fetch-depth: 0 + + - name: Mark repository as safe + shell: bash + run: git config --global --add safe.directory "$GITHUB_WORKSPACE" + + - name: Validate backfill target + if: github.event_name == 'workflow_dispatch' + shell: bash + env: + RELEASE_TAG: ${{ inputs.tag }} + run: | + set -euo pipefail + if [[ "${RELEASE_TAG}" != v* || "${RELEASE_TAG}" == *-* ]]; then + echo "::error::Backfill tag must be a stable release tag starting with v" + exit 1 + fi + git check-ref-format "refs/tags/${RELEASE_TAG}" + TAG_SHA="$(git rev-parse --verify "refs/tags/${RELEASE_TAG}^{commit}")" + BUILD_SHA="$(git rev-parse HEAD)" + if [[ "${BUILD_SHA}" != "${TAG_SHA}" ]]; then + echo "::error::Build source ${BUILD_SHA} does not match ${RELEASE_TAG} at ${TAG_SHA}" + exit 1 + fi + echo "Build source: ${RELEASE_TAG} at ${BUILD_SHA}" + + - name: Setup Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Setup Go + uses: actions/setup-go@v6 + with: + go-version: '1.25.9' + cache: false + + - name: Set Python bin + shell: bash + run: | + set -euo pipefail + PYTHON_BIN="$(command -v python${{ matrix.python-version }})" + "$PYTHON_BIN" --version + "$PYTHON_BIN" -m pip --version + echo "PYTHON_BIN=${PYTHON_BIN}" >> "$GITHUB_ENV" + + - name: Install dependencies + shell: bash + run: | + set -eo pipefail + export PYTHONPATH="${PYTHONPATH:-}" + export CMAKE_PREFIX_PATH="${CMAKE_PREFIX_PATH:-}" + bash -x dependencies.sh -y + echo "PATH=/usr/local/musa/bin:${PATH}" >> "$GITHUB_ENV" + echo "LD_LIBRARY_PATH=/usr/local/musa/lib:${LD_LIBRARY_PATH:-}" >> "$GITHUB_ENV" + echo "LIBRARY_PATH=/usr/local/musa/lib:${LIBRARY_PATH:-}" >> "$GITHUB_ENV" + + - name: Verify Go toolchain + shell: bash + run: | + set -euo pipefail + GO_BIN="$(command -v go || true)" + if [[ -z "${GO_BIN}" ]]; then + echo "::error::Go is not available on PATH" + exit 1 + fi + echo "Go executable: ${GO_BIN}" + "${GO_BIN}" version + + - name: Configure project + shell: bash + run: | + set -eo pipefail + export PYTHONPATH="${PYTHONPATH:-}" + export CMAKE_PREFIX_PATH="${CMAKE_PREFIX_PATH:-}" + rm -rf build + mkdir build + cd build + cmake_args=( + -DUSE_MUSA=ON + -DUSE_HTTP=ON + -DUSE_ETCD=ON + -DSTORE_USE_ETCD=ON + -DBUILD_UNIT_TESTS=OFF + -DENABLE_DEBUG_SYMBOLS=OFF + -DCMAKE_BUILD_TYPE=Release + -DPython3_EXECUTABLE="${PYTHON_BIN}" + ) + cmake -G Ninja .. "${cmake_args[@]}" + + - name: Build project + shell: bash + run: | + set -eo pipefail + export PYTHONPATH="${PYTHONPATH:-}" + export CMAKE_PREFIX_PATH="${CMAKE_PREFIX_PATH:-}" + cd build + cmake --build . -j"$(nproc)" + + - name: Build MUSA allocator + shell: bash + run: | + set -eo pipefail + mkdir -p build/mooncake-transfer-engine/nvlink-allocator + cd mooncake-transfer-engine/nvlink-allocator + bash build.sh --use-mcc ../../build/mooncake-transfer-engine/nvlink-allocator/ + + - name: Install project + shell: bash + run: | + set -eo pipefail + cd build + cmake --install . + + - name: Generate Python version tag + id: generate_tag_release + shell: bash + run: | + echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT + + - name: Build Python wheel + shell: bash + run: | + set -eo pipefail + export LD_LIBRARY_PATH="${LD_LIBRARY_PATH:-}:/usr/local/lib:/usr/local/musa/lib" + MUSA_BUILD=1 PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-musa-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh + + - name: Upload Python wheel artifact + uses: actions/upload-artifact@v4 + with: + name: mooncake-wheel-musa-x86_64-py${{ steps.generate_tag_release.outputs.python_version_tag }} + path: mooncake-wheel/dist-musa-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl + + publish-release: + if: ${{ github.event_name == 'workflow_dispatch' || !contains(github.ref_name, '-') }} + needs: build + runs-on: ubuntu-22.04 + environment: pypi + + permissions: + contents: write + id-token: write + + steps: + - name: Checkout source + uses: actions/checkout@v4 + with: + ref: ${{ github.event_name == 'workflow_dispatch' && format('refs/tags/{0}', inputs.tag) || github.ref }} + + - name: Download all wheel artifacts + uses: actions/download-artifact@v4 + with: + path: mooncake-wheel/dist-all + pattern: mooncake-wheel-musa-* + + - name: Prepare wheels for release + run: | + mkdir -p mooncake-wheel/dist-release + find mooncake-wheel/dist-all -name "*.whl" -exec cp {} mooncake-wheel/dist-release/ \; + echo "Collected wheels for release:" + ls -la mooncake-wheel/dist-release/ + + - name: Upload wheels to GitHub Release + uses: softprops/action-gh-release@v1 + with: + tag_name: ${{ github.event_name == 'workflow_dispatch' && inputs.tag || github.ref_name }} + files: mooncake-wheel/dist-release/*.whl + + - name: Publish package to PyPI + if: github.repository == 'kvcache-ai/Mooncake' + uses: pypa/gh-action-pypi-publish@release/v1 + with: + packages-dir: mooncake-wheel/dist-release/ diff --git a/.github/workflows/release-non-cuda.yaml b/.github/workflows/release-non-cuda.yaml index ab9b84243c..7d52069fb2 100644 --- a/.github/workflows/release-non-cuda.yaml +++ b/.github/workflows/release-non-cuda.yaml @@ -5,119 +5,33 @@ on: tags: - 'v*' -env: - SCCACHE_GHA_ENABLED: "true" jobs: + # manylinux2_28 for the toolchain only (USE_CUDA=OFF); keeps the glibc floor + # aligned with the CUDA wheels. build: - runs-on: ubuntu-22.04 - permissions: - contents: write + if: ${{ !contains(github.ref_name, '-') }} strategy: + fail-fast: false matrix: - python-version: ['3.10', '3.11', '3.12', '3.13'] - env: - BUILD_WITH_EP: "0" - NON_CUDA_BUILD: "1" - steps: - - name: Checkout source - uses: actions/checkout@v4 - - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - - - name: Free up disk space - run: | - sudo rm -rf /usr/share/dotnet - sudo rm -rf /opt/ghc - sudo rm -rf /opt/hostedtoolcache/CodeQL - - - name: Run sccache-cache - uses: mozilla-actions/sccache-action@v0.0.9 - - - name: Configure sccache - uses: actions/github-script@v7 - with: - script: | - core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); - core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); - - - name: Run sccache stat for check - shell: bash - run: ${SCCACHE_PATH} --show-stats - - - name: Configure project - run: | - sudo apt update -y - sudo bash -x dependencies.sh -y - mkdir build - cd build - cmake .. -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=OFF -DWITH_EP=OFF -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release - shell: bash - - - name: Build project - run: | - cd build - make -j - sudo make install - shell: bash - - - name: Generate Python version tag - id: generate_tag_release - run: | - echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT - shell: bash - - - name: Build Python wheel - run: | - # Set LD_LIBRARY_PATH for wheel building - export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib - PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh - env: - VERSION: ${{ env.VERSION }} - - - name: Upload Python wheel artifact - uses: actions/upload-artifact@v4 - with: - name: mooncake-wheel-non-cuda-py${{ steps.generate_tag_release.outputs.python_version_tag }} - path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl + include: + - architecture: x86_64 + artifact-prefix: mooncake-wheel-non-cuda + - architecture: arm64 + artifact-prefix: mooncake-wheel-non-cuda-arm64 + uses: ./.github/workflows/_build-wheel.yaml + with: + variant: non-cuda + architecture: ${{ matrix.architecture }} + artifact-prefix: ${{ matrix.artifact-prefix }} publish-release: + if: ${{ !contains(github.ref_name, '-') }} needs: build - runs-on: ubuntu-22.04 permissions: contents: write id-token: write - steps: - - name: Checkout source - uses: actions/checkout@v4 - - - name: Download all wheel artifacts - uses: actions/download-artifact@v4 - with: - path: mooncake-wheel/dist-all - pattern: mooncake-wheel-non-cuda-py* - - - name: Prepare wheels for release - run: | - # Move all wheels to a single directory - mkdir -p mooncake-wheel/dist-release - find mooncake-wheel/dist-all -name "*.whl" -exec cp {} mooncake-wheel/dist-release/ \; - ls -la mooncake-wheel/dist-release/ - # List all collected wheels - echo "Collected wheels for release:" - ls -la mooncake-wheel/dist-release/ - - - name: Upload wheels to GitHub Release - uses: softprops/action-gh-release@v1 - with: - files: mooncake-wheel/dist-release/*.whl - - - name: Publish package to PyPI - if: github.repository == 'kvcache-ai/Mooncake' - uses: pypa/gh-action-pypi-publish@release/v1 - with: - packages-dir: mooncake-wheel/dist-release/ - password: ${{ secrets.PYPI_API_TOKEN }} - + uses: ./.github/workflows/_publish-wheel.yaml + with: + artifact-pattern: 'mooncake-wheel-non-cuda-*' + secrets: + pypi-token: ${{ secrets.PYPI_API_TOKEN }} diff --git a/.github/workflows/release-npu.yaml b/.github/workflows/release-npu.yaml new file mode 100644 index 0000000000..9ef99593d3 --- /dev/null +++ b/.github/workflows/release-npu.yaml @@ -0,0 +1,195 @@ +name: Release Ascend NPU + +on: + push: + tags: + - 'v*' + workflow_dispatch: + inputs: + tag: + description: 'Release tag (e.g. v0.1.0)' + required: true + type: string + +jobs: + build: + if: ${{ (github.event_name == 'push' && !contains(github.ref_name, '-')) || (github.event_name == 'workflow_dispatch' && startsWith(inputs.tag, 'v') && !contains(inputs.tag, '-')) }} + + strategy: + max-parallel: 2 + matrix: + arch: [aarch64, x86_64] + python-version: ['3.9', '3.10', '3.11', '3.12', '3.13'] + include: + - arch: aarch64 + runner: ubuntu-22.04-arm + cann_arch: aarch64 + - arch: x86_64 + runner: ubuntu-22.04 + cann_arch: x86_64 + + runs-on: ${{ matrix.runner }} + + permissions: + contents: write + + env: + NPU_BUILD: "1" + + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ github.event_name == 'workflow_dispatch' && format('refs/tags/{0}', inputs.tag) || github.ref }} + fetch-depth: 0 + + - name: Setup Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Set Python bin + shell: bash + run: | + set -euo pipefail + PYTHON_BIN="$(command -v python${{ matrix.python-version }})" + "$PYTHON_BIN" --version + "$PYTHON_BIN" -m pip --version + echo "PYTHON_BIN=${PYTHON_BIN}" >> "$GITHUB_ENV" + + - name: Install CANN Toolkit + shell: bash + run: | + set -euo pipefail + CANN_BASE_URL="https://ascend.devcloud.huaweicloud.com/cann/run/software" + CANN_VERSION=$(curl -s "${CANN_BASE_URL}/" | grep -oP '[0-9]+\.[0-9]+\.[0-9]+(?=/)' | sort -V | tail -1) + echo "Latest CANN version: ${CANN_VERSION}" + CANN_URL="${CANN_BASE_URL}/${CANN_VERSION}/${{ matrix.cann_arch }}/Ascend-cann-toolkit_${CANN_VERSION}_linux-${{ matrix.cann_arch }}.run" + echo "Downloading CANN ${CANN_VERSION} from ${CANN_URL}" + wget -q --show-progress -O /tmp/cann_toolkit.run "${CANN_URL}" + chmod +x /tmp/cann_toolkit.run + sudo /tmp/cann_toolkit.run --install --install-for-all --install-path=/usr/local/Ascend --quiet + rm -f /tmp/cann_toolkit.run + + - name: Build and install HIXL + shell: bash + run: | + set -eo pipefail + export PYTHONPATH="${PYTHONPATH:-}" + export CMAKE_PREFIX_PATH="${CMAKE_PREFIX_PATH:-}" + source /usr/local/Ascend/cann/set_env.sh + sudo ${PYTHON_BIN} -m pip install setuptools + cd /tmp + git clone https://gitcode.com/cann/hixl.git + cd hixl + bash build.sh -j$(nproc) + sudo ./build_out/cann-hixl_*.run --full --quiet --pylocal --install-path=/usr/local/Ascend + + - name: Install dependencies + shell: bash + run: | + set -eo pipefail + export PYTHONPATH="${PYTHONPATH:-}" + export CMAKE_PREFIX_PATH="${CMAKE_PREFIX_PATH:-}" + source /usr/local/Ascend/cann/set_env.sh + sudo bash -x dependencies.sh -y + sudo bash scripts/ascend/dependencies_ascend_installation.sh + echo "PATH=/usr/local/go/bin:${PATH}" >> "$GITHUB_ENV" + + - name: Configure project + shell: bash + run: | + set -eo pipefail + export PYTHONPATH="${PYTHONPATH:-}" + export CMAKE_PREFIX_PATH="${CMAKE_PREFIX_PATH:-}" + source /usr/local/Ascend/cann/set_env.sh + rm -rf build + mkdir build + cd build + cmake_args=( + -DUSE_ASCEND_DIRECT=ON + -DUSE_ETCD=ON + -DSTORE_USE_ETCD=ON + -DBUILD_UNIT_TESTS=OFF + -DCMAKE_BUILD_TYPE=Release + -DPython3_EXECUTABLE="${PYTHON_BIN}" + -DCMAKE_INSTALL_RPATH_USE_LINK_PATH=OFF + -DCMAKE_BUILD_WITH_INSTALL_RPATH=OFF + -DCMAKE_SKIP_INSTALL_RPATH=ON + ) + cmake .. "${cmake_args[@]}" + + - name: Build project + shell: bash + run: | + set -eo pipefail + export PYTHONPATH="${PYTHONPATH:-}" + export CMAKE_PREFIX_PATH="${CMAKE_PREFIX_PATH:-}" + source /usr/local/Ascend/cann/set_env.sh + export PATH="/usr/local/go/bin:${PATH}" + cd build + cmake --build . -j"$(nproc)" + sudo cmake --install . + + - name: Generate Python version tag + id: generate_tag_release + shell: bash + run: | + echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT + + - name: Build Python wheel + shell: bash + run: | + set -eo pipefail + export PYTHONPATH="${PYTHONPATH:-}" + export CMAKE_PREFIX_PATH="${CMAKE_PREFIX_PATH:-}" + source /usr/local/Ascend/cann/set_env.sh + export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib + NPU_BUILD=1 PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-npu-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh + + - name: Upload Python wheel artifact + uses: actions/upload-artifact@v4 + with: + name: mooncake-wheel-npu-${{ matrix.arch }}-py${{ steps.generate_tag_release.outputs.python_version_tag }} + path: mooncake-wheel/dist-npu-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl + + publish-release: + if: ${{ (github.event_name == 'push' && !contains(github.ref_name, '-')) || (github.event_name == 'workflow_dispatch' && startsWith(inputs.tag, 'v') && !contains(inputs.tag, '-')) }} + needs: build + runs-on: ubuntu-22.04 + environment: pypi + + permissions: + contents: write + id-token: write + + steps: + - name: Checkout source + uses: actions/checkout@v4 + with: + ref: ${{ github.event_name == 'workflow_dispatch' && format('refs/tags/{0}', inputs.tag) || github.ref }} + + - name: Download all wheel artifacts + uses: actions/download-artifact@v4 + with: + path: mooncake-wheel/dist-all + pattern: mooncake-wheel-npu-* + + - name: Prepare wheels for release + run: | + mkdir -p mooncake-wheel/dist-release + find mooncake-wheel/dist-all -name "*.whl" -exec cp {} mooncake-wheel/dist-release/ \; + echo "Collected wheels for release:" + ls -la mooncake-wheel/dist-release/ + + - name: Upload wheels to GitHub Release + uses: softprops/action-gh-release@v1 + with: + tag_name: ${{ github.event_name == 'workflow_dispatch' && inputs.tag || github.ref_name }} + files: mooncake-wheel/dist-release/*.whl + + - name: Publish package to PyPI + if: github.repository == 'kvcache-ai/Mooncake' + uses: pypa/gh-action-pypi-publish@release/v1 + with: + packages-dir: mooncake-wheel/dist-release/ diff --git a/.github/workflows/release-rocm.yaml b/.github/workflows/release-rocm.yaml new file mode 100644 index 0000000000..060285da36 --- /dev/null +++ b/.github/workflows/release-rocm.yaml @@ -0,0 +1,272 @@ +name: Release ROCm + +on: + push: + tags: + - 'v*' + workflow_dispatch: + inputs: + tag: + # Only tags whose tree already contains ROCm support (the HIP_BUILD + # variant in scripts/build_wheel.sh, i.e. this PR onward) can be built. + # Tags predating ROCm support are rejected by the guard below, since + # they would produce a mislabeled default-named wheel. + description: 'Release tag to (re)build the ROCm wheel for (must contain HIP_BUILD support)' + required: true + type: string + +concurrency: + group: release-rocm-${{ github.event_name == 'workflow_dispatch' && inputs.tag || github.ref_name }} + cancel-in-progress: false + +jobs: + build: + if: ${{ github.event_name == 'workflow_dispatch' || !contains(github.ref_name, '-') }} + runs-on: ubuntu-22.04 + # ROCm 7.2 dev image: matches the upstream vllm/vllm-openai-rocm and sglang + # ROCm images. hipcc/HIP/hsa-runtime are present; no GPU needed to compile. + container: rocm/dev-ubuntu-22.04:7.2.3-complete + + permissions: + contents: write + + strategy: + max-parallel: 2 + matrix: + python-version: ['3.10', '3.11', '3.12', '3.13'] + + env: + HIP_BUILD: "1" + + steps: + # git must exist BEFORE checkout so actions/checkout does a real clone + # (with .git + submodules); the rocm/dev image ships without git. + - name: Install git (pre-checkout) + shell: bash + run: | + set -eo pipefail + export DEBIAN_FRONTEND=noninteractive + apt-get update -y + apt-get install -y --no-install-recommends git ca-certificates + + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ github.event_name == 'workflow_dispatch' && format('refs/tags/{0}', inputs.tag) || github.ref }} + fetch-depth: 0 + submodules: recursive + + - name: Mark repository as safe + shell: bash + run: git config --global --add safe.directory "$GITHUB_WORKSPACE" || true + + - name: Validate backfill target + if: github.event_name == 'workflow_dispatch' + shell: bash + env: + RELEASE_TAG: ${{ inputs.tag }} + run: | + set -euo pipefail + if [[ "${RELEASE_TAG}" != v* || "${RELEASE_TAG}" == *-* ]]; then + echo "::error::Backfill tag must be a stable release tag starting with v" + exit 1 + fi + git check-ref-format "refs/tags/${RELEASE_TAG}" + TAG_SHA="$(git rev-parse --verify "refs/tags/${RELEASE_TAG}^{commit}")" + BUILD_SHA="$(git rev-parse HEAD)" + if [[ "${BUILD_SHA}" != "${TAG_SHA}" ]]; then + echo "::error::Build source ${BUILD_SHA} does not match ${RELEASE_TAG} at ${TAG_SHA}" + exit 1 + fi + + - name: Ensure checked-out tree supports the ROCm wheel variant + shell: bash + run: | + set -eo pipefail + # Guard against backfilling a tag that predates ROCm support: its + # scripts/build_wheel.sh would ignore HIP_BUILD and emit the default + # "mooncake-transfer-engine" package, which the publish steps would + # then upload over the CUDA package on PyPI. Refuse such tags. + if ! grep -q 'HIP_BUILD' scripts/build_wheel.sh; then + echo "::error::Checked-out tree has no HIP_BUILD support in scripts/build_wheel.sh;" + echo "::error::refusing to build a mislabeled ROCm wheel. Use a tag at or after ROCm support landed." + exit 1 + fi + + - name: Install toolchain and Python ${{ matrix.python-version }} + shell: bash + run: | + set -eo pipefail + export DEBIAN_FRONTEND=noninteractive + apt-get update -y + apt-get install -y --no-install-recommends \ + git curl ca-certificates build-essential sudo pkg-config \ + ninja-build software-properties-common + PYV="${{ matrix.python-version }}" + if ! command -v "python${PYV}" >/dev/null 2>&1; then + add-apt-repository -y ppa:deadsnakes/ppa + apt-get update -y + fi + # Always install -dev + -venv: the image's system python3.10 exists but + # ships without the venv module / dev headers. + apt-get install -y --no-install-recommends \ + "python${PYV}" "python${PYV}-dev" "python${PYV}-venv" + curl -sS https://bootstrap.pypa.io/get-pip.py | "python${PYV}" + PYTHON_BIN="$(command -v python${PYV})" + echo "PYTHON_BIN=${PYTHON_BIN}" >> "$GITHUB_ENV" + echo "/opt/rocm/bin" >> "$GITHUB_PATH" + + - name: Install dependencies + shell: bash + run: | + set -eo pipefail + bash -x dependencies.sh -y + echo "/usr/local/go/bin" >> "$GITHUB_PATH" + + - name: Configure project + shell: bash + run: | + set -eo pipefail + export PATH="/opt/rocm/bin:$PATH" # hipify-perl for USE_HIP + rm -rf build && mkdir build && cd build + cmake -G Ninja .. \ + -DUSE_HIP=ON -DUSE_CUDA=OFF -DWITH_EP=OFF \ + -DUSE_HTTP=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON \ + -DBUILD_UNIT_TESTS=OFF -DENABLE_DEBUG_SYMBOLS=OFF \ + -DCMAKE_BUILD_TYPE=Release \ + -DPython3_EXECUTABLE="${PYTHON_BIN}" + + - name: Build project + shell: bash + run: | + set -eo pipefail + export PATH="/opt/rocm/bin:$PATH" + cd build + # Retry to ride out transient Go module (proxy.golang.org) fetch errors + # during the etcd-wrapper build; Go caches modules, so retries resume. + n=0 + until cmake --build . -j"$(nproc)"; do + n=$((n+1)) + if [ "$n" -ge 3 ]; then echo "Build failed after $n attempts"; exit 1; fi + echo "Build attempt $n failed; retrying in 15s..."; sleep 15 + done + cmake --install . + + - name: Generate Python version tag + id: generate_tag_release + shell: bash + run: echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> "$GITHUB_OUTPUT" + + - name: Build Python wheel + shell: bash + run: | + set -eo pipefail + export PATH="/opt/rocm/bin:$PATH" + export LD_LIBRARY_PATH="${LD_LIBRARY_PATH:-}:/usr/local/lib" + HIP_BUILD=1 PYTHON_VERSION=${{ matrix.python-version }} \ + OUTPUT_DIR=dist-rocm-py${{ steps.generate_tag_release.outputs.python_version_tag }} \ + ./scripts/build_wheel.sh + + - name: Assert ROCm package name + shell: bash + run: | + set -eo pipefail + # Defense in depth: never let a wheel that was not renamed to the ROCm + # package (e.g. HIP_BUILD silently ignored) reach the publish steps. + dir="mooncake-wheel/dist-rocm-py${{ steps.generate_tag_release.outputs.python_version_tag }}" + if ! ls "${dir}"/mooncake_transfer_engine_rocm-*.whl >/dev/null 2>&1; then + echo "::error::Built wheel is not named mooncake-transfer-engine-rocm:" + ls -la "${dir}" || true + exit 1 + fi + + - name: Smoke test the exact release wheel + shell: bash + run: | + set -eo pipefail + # Install the actual per-matrix release artifact (cp310-cp313) into a + # fresh venv and exercise the packaged binary, so packaging/ELF-layout + # failures in any version are caught before this wheel reaches PyPI. + # No GPU is needed for `mooncake_master --version`. + smoke_venv=$(mktemp -d) + "${PYTHON_BIN}" -m venv "$smoke_venv" + "$smoke_venv/bin/python" -m pip install --no-deps \ + mooncake-wheel/dist-rocm-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl + export LD_LIBRARY_PATH="/opt/rocm/lib:/usr/local/lib:${LD_LIBRARY_PATH:-}" + site="$("$smoke_venv/bin/python" -c 'import mooncake,os;print(os.path.dirname(mooncake.__file__))')" + "$site/mooncake_master" --version + + - name: Upload Python wheel artifact + uses: actions/upload-artifact@v4 + with: + name: mooncake-wheel-rocm-x86_64-py${{ steps.generate_tag_release.outputs.python_version_tag }} + path: mooncake-wheel/dist-rocm-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl + + # Publish to PyPI and attach to the GitHub Release in SEPARATE jobs, with the + # GitHub Release job depending on PyPI. PyPI uploads are immutable, so if a + # single combined job published to PyPI and then failed on the GitHub Release + # upload, "re-run failed jobs" would restart at the (now-rejected) PyPI step + # and could never reach the release upload. Split jobs let a transient GitHub + # Release failure be retried on its own without republishing PyPI files. + publish-pypi: + if: ${{ github.event_name == 'workflow_dispatch' || !contains(github.ref_name, '-') }} + needs: build + runs-on: ubuntu-22.04 + environment: pypi + + permissions: + id-token: write + + steps: + - name: Download all wheel artifacts + uses: actions/download-artifact@v4 + with: + path: mooncake-wheel/dist-all + pattern: mooncake-wheel-rocm-* + + - name: Prepare wheels for release + run: | + mkdir -p mooncake-wheel/dist-release + find mooncake-wheel/dist-all -name "*.whl" -exec cp {} mooncake-wheel/dist-release/ \; + echo "Collected wheels for release:" + ls -la mooncake-wheel/dist-release/ + + # PREREQUISITE (one-time, maintainer/PyPI-owner action): this uses OIDC + # (no API token), so PyPI must have a Trusted Publisher configured for + # the NEW project `mooncake-transfer-engine-rocm` with identity + # owner/repo `kvcache-ai/Mooncake`, workflow `release-rocm.yaml`, and + # environment `pypi`. Without it the publish step fails. + - name: Publish package to PyPI + if: github.repository == 'kvcache-ai/Mooncake' + uses: pypa/gh-action-pypi-publish@release/v1 + with: + packages-dir: mooncake-wheel/dist-release/ + + publish-github-release: + # Runs only after PyPI succeeds; can be re-run on its own if the GitHub + # Release upload fails transiently, without touching the immutable PyPI files. + if: ${{ github.event_name == 'workflow_dispatch' || !contains(github.ref_name, '-') }} + needs: publish-pypi + runs-on: ubuntu-22.04 + + permissions: + contents: write + + steps: + - name: Download all wheel artifacts + uses: actions/download-artifact@v4 + with: + path: mooncake-wheel/dist-all + pattern: mooncake-wheel-rocm-* + + - name: Prepare wheels for release + run: | + mkdir -p mooncake-wheel/dist-release + find mooncake-wheel/dist-all -name "*.whl" -exec cp {} mooncake-wheel/dist-release/ \; + ls -la mooncake-wheel/dist-release/ + + - name: Upload wheels to GitHub Release + uses: softprops/action-gh-release@v1 + with: + tag_name: ${{ github.event_name == 'workflow_dispatch' && inputs.tag || github.ref_name }} + files: mooncake-wheel/dist-release/*.whl diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index e37e72838f..f5cb1c5ff4 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -5,139 +5,30 @@ on: tags: - 'v*' -env: - SCCACHE_GHA_ENABLED: "true" jobs: + # Skip semver pre-release tags (e.g. v1.0.0-rc1); those are handled by pre-release.yaml. build: - runs-on: ubuntu-22.04 - permissions: - contents: write + if: ${{ !contains(github.ref_name, '-') }} strategy: + fail-fast: false matrix: - python-version: ['3.10', '3.11', '3.12', '3.13'] - env: - BUILD_WITH_EP: "1" - TORCH_CUDA_ARCH_LIST: "8.0;9.0" - steps: - - name: Checkout source - uses: actions/checkout@v4 - - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - - - name: Free up disk space - run: | - sudo rm -rf /usr/share/dotnet - sudo rm -rf /opt/ghc - sudo rm -rf /opt/hostedtoolcache/CodeQL - sudo rm -rf /usr/local/lib/android - df -h - - - name: Install CUDA Toolkit - uses: Jimver/cuda-toolkit@v0.2.24 - with: - cuda: '12.8.1' - linux-local-args: '["--toolkit"]' - method: 'network' - sub-packages: '["nvcc", "nvrtc-dev"]' - non-cuda-sub-packages: '["libcusparse-dev", "libcublas-dev", "libcusolver-dev"]' - - - name: Run sccache-cache - uses: mozilla-actions/sccache-action@v0.0.9 - - - name: Configure sccache - uses: actions/github-script@v7 - with: - script: | - core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); - core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); - - - name: Run sccache stat for check - shell: bash - run: ${SCCACHE_PATH} --show-stats - - - name: Configure project - run: | - sudo apt update -y - 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" -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release - shell: bash - - - name: Build project - run: | - export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:LIBRARY_PATH - cd build - make -j - sudo make install - shell: bash - - - name: Build nvlink_allocator.so - run: | - export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH - export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH - export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH - mkdir -p build/mooncake-transfer-engine/nvlink-allocator - cd mooncake-transfer-engine/nvlink-allocator - bash build.sh ../../build/mooncake-transfer-engine/nvlink-allocator/ - shell: bash - - - name: Generate Python version tag - id: generate_tag_release - run: | - echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT - shell: bash - - - name: Build Python wheel - run: | - # Set LD_LIBRARY_PATH for wheel building - export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib - PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh - env: - VERSION: ${{ env.VERSION }} - - - name: Upload Python wheel artifact - uses: actions/upload-artifact@v4 - with: - name: mooncake-wheel-py${{ steps.generate_tag_release.outputs.python_version_tag }} - path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl + include: + - architecture: x86_64 + artifact-prefix: mooncake-wheel + - architecture: arm64 + artifact-prefix: mooncake-wheel-arm64 + uses: ./.github/workflows/_build-wheel.yaml + with: + architecture: ${{ matrix.architecture }} + artifact-prefix: ${{ matrix.artifact-prefix }} publish-release: needs: build - runs-on: ubuntu-22.04 permissions: contents: write id-token: write - steps: - - name: Checkout source - uses: actions/checkout@v4 - - - name: Download all wheel artifacts - uses: actions/download-artifact@v4 - with: - path: mooncake-wheel/dist-all - - - name: Prepare wheels for release - run: | - # Move all wheels to a single directory - mkdir -p mooncake-wheel/dist-release - find mooncake-wheel/dist-all -name "*.whl" -exec cp {} mooncake-wheel/dist-release/ \; - ls -la mooncake-wheel/dist-release/ - # List all collected wheels - echo "Collected wheels for release:" - ls -la mooncake-wheel/dist-release/ - - - name: Upload wheels to GitHub Release - uses: softprops/action-gh-release@v1 - with: - files: mooncake-wheel/dist-release/*.whl - - - name: Publish package to PyPI - if: github.repository == 'kvcache-ai/Mooncake' - uses: pypa/gh-action-pypi-publish@release/v1 - with: - packages-dir: mooncake-wheel/dist-release/ - password: ${{ secrets.PYPI_API_TOKEN }} + uses: ./.github/workflows/_publish-wheel.yaml + with: + artifact-pattern: 'mooncake-wheel*' + secrets: + pypi-token: ${{ secrets.PYPI_API_TOKEN }} diff --git a/.github/workflows/stale.yml b/.github/workflows/stale.yml new file mode 100644 index 0000000000..d4b92c8f96 --- /dev/null +++ b/.github/workflows/stale.yml @@ -0,0 +1,41 @@ +name: Close Stale Issues and PRs + +on: + schedule: + - cron: '37 2 * * *' + workflow_dispatch: {} + +permissions: + issues: write + pull-requests: write + +jobs: + stale: + runs-on: ubuntu-latest + steps: + - uses: actions/stale@v9 + with: + repo-token: ${{ secrets.GITHUB_TOKEN }} + stale-issue-message: > + This issue has been automatically marked as stale because it has not had + recent activity. It will be closed in 30 days if no further activity occurs. + If this issue is still relevant, please comment to keep it open. + stale-pr-message: > + This PR has been automatically marked as stale because it has not had + recent activity. It will be closed in 14 days if no further activity occurs. + Please update the PR or comment to keep it open. + close-issue-message: > + This issue was closed because it has been inactive for 120 days. + Feel free to reopen if this is still relevant. + close-pr-message: > + This PR was closed because it has been inactive for 104 days. + Feel free to reopen if you'd like to continue the work. + days-before-stale: 120 + days-before-close: 30 + days-before-pr-stale: 120 + days-before-pr-close: 30 + stale-issue-label: 'stale' + stale-pr-label: 'stale' + exempt-issue-labels: 'pinned,security,RFC' + exempt-pr-labels: 'pinned,work-in-progress' + operations-per-run: 100 diff --git a/.gitignore b/.gitignore index 23b51fadbb..35a253a732 100644 --- a/.gitignore +++ b/.gitignore @@ -5,6 +5,7 @@ build_ofed4 old local_test go.sum +!mooncake-common/etcd/go.sum *.so bin mod @@ -197,11 +198,14 @@ mooncake-wheel/mooncake/allocator_ascend_npu.py mooncake-wheel/mooncake/mooncake_master mooncake-wheel/mooncake/transfer_engine_bench -# Claude Code Memory -CLAUDE.md - # CodeQL _codeql_detected_source_root # CodeBuddy Memory .codebuddy/ + +# core dumps +core_* +# MacOS +.DS_Store +.envrc diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index c8f2846cb6..06600a0e65 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,6 +1,9 @@ # Pre-commit hooks configuration for Mooncake # Install: pip install -r requirements-dev.txt && pre-commit install -# Run manually: pre-commit run --all-files +# Staged files: pre-commit run +# PR-changed files: git fetch origin main && pre-commit run --files $(git diff --name-only --diff-filter=ACMR origin/main...HEAD) +# Full-repo (intentional cleanup only): pre-commit run --all-files +# Format all C/C++ files explicitly: ./scripts/code_format.sh --all # Note: clang-format should already be available (installed via system packages or dependencies.sh) # Exclusions: build artifacts, vendored extern code, generated wheels. @@ -9,6 +12,9 @@ minimum_pre_commit_version: '3.6.0' ci: autofix_prs: true autoupdate_commit_msg: 'chore: pre-commit autoupdate' + # Needs Cargo + libclang, which the hosted pre-commit.ci image lacks; + # GitHub Actions is the authoritative check for this hook. + skip: [store-rust-dlopen-bindings] repos: - repo: https://github.com/pre-commit/pre-commit-hooks @@ -23,6 +29,25 @@ repos: - id: check-added-large-files args: ['--maxkb=1024'] + - repo: local + hooks: + - id: mooncake-code-format + name: Format staged C/C++ changes + entry: ./scripts/code_format.sh --staged + language: system + files: '\.(c|cc|cpp|cxx|cu|cuh|h|hpp)$' + exclude: '^(extern/|build/|.*/build/|FAST25-release/|.*cachelib_memory_allocator/|.*/thirdparty/.*)' + require_serial: true + # Regenerate the committed dlopen bindings when store_c.h or the generator + # inputs change; a resulting diff fails the commit so they can't go stale. + - id: store-rust-dlopen-bindings + name: Regenerate Mooncake Store Rust dlopen bindings + entry: cargo run --locked --manifest-path mooncake-store/rust/Cargo.toml --example generate_dlopen_bindings + language: system + pass_filenames: false + require_serial: true + files: '^(mooncake-store/include/store_c\.h|mooncake-store/rust/(Cargo\.(toml|lock)|examples/generate_dlopen_bindings\.rs|src/generated/ffi_dlopen_bindings\.rs))$' + - repo: https://github.com/astral-sh/ruff-pre-commit rev: v0.6.9 hooks: @@ -37,14 +62,7 @@ repos: hooks: - id: codespell exclude: '^(extern/|FAST25-release/)' - args: ['--ignore-words-list=te,mooncake,KVCache'] - - - repo: https://github.com/pre-commit/mirrors-clang-format - rev: v20.1.8 - hooks: - - id: clang-format - files: '\.(c|cc|cpp|cxx|h|hpp)$' - exclude: '^(extern/|build/|.*/build/|FAST25-release/|.*/thirdparty/.*)' + args: ['--ignore-words-list=te,mooncake,KVCache,cann,hsa,crate'] - repo: https://github.com/cheshirekow/cmake-format-precommit rev: v0.6.13 diff --git a/.typos.toml b/.typos.toml index aa808bb237..4843e506f2 100644 --- a/.typos.toml +++ b/.typos.toml @@ -1,12 +1,21 @@ [default] -extend-ignore-words = ["CANN", "ASO", "fre"] +extend-ignore-words = ["CANN", "ASO", "fre", "wqs", "hsa", "ue"] [default.extend-words] CANN = "CANN" ASO = "ASO" fre = "fre" +wqs = "wqs" +# AMD HSA runtime symbol prefix (hsa_*, hsaRes, hsaErr, etc.) — used by the +# ROCm dmabuf MR registration path. +hsa = "hsa" +Optin = "Optin" +HPE = "HPE" [files] extend-exclude = [ "mooncake-transfer-engine/tent/include/tent/thirdparty/nlohmann/json.h", -] \ No newline at end of file + # DeepEP-derived elastic kernel headers keep upstream identifiers such as + # `ue8m0x4`; exclude the imported header block from spelling checks. + "mooncake-ep/include/elastic/*", +] diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000000..6ab75b16b9 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,27 @@ +# AGENTS.md + +## `docs/` Directory Changes + +- Before modifying files under `docs/`, read `docs/AGENTS.md`. + +## Pull Request Guidelines + +- Follow `CONTRIBUTING.md` for PR title prefixes, RFC expectations, and + contribution workflow. +- Before opening a PR for nontrivial work, check whether an existing issue or + open PR already covers the same change. If the work overlaps, explain the + difference instead of duplicating it. +- Do not open low-value busywork PRs for isolated typo, style, or mechanical + changes unless they are part of a substantive requested change. +- Use `.github/pull_request_template.md` when preparing a PR, and fill in the + relevant sections for description, module, type of change, testing, + checklist, and AI assistance disclosure. +- For AI-assisted changes, make sure the human submitter has reviewed every + changed line and can defend the change end-to-end. +- Before handoff, run pre-commit on the files touched by the change when the + toolchain is available (see `CONTRIBUTING.md` for the PR-scoped + `pre-commit run --files ...` command). Do not use + `pre-commit run --all-files` for routine PRs; if it rewrites unrelated + files, leave those edits out of the PR. +- Keep PRs lean: review `git diff` before staging, and include only changes + required for the requested task. diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000000..43c994c2d3 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1 @@ +@AGENTS.md diff --git a/CMakeLists.txt b/CMakeLists.txt index ec9d7ef9d0..a3cf8ce07e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -4,131 +4,276 @@ project(mooncake CXX C) # indicates cmake is invoked from top-level dir set(GLOBAL_CONFIG "true") +if(SKBUILD) + add_subdirectory(python) +endif() + include(mooncake-common/FindJsonCpp.cmake) include(mooncake-common/FindGLOG.cmake) include(mooncake-common/common.cmake) # unit test -if (BUILD_UNIT_TESTS) +if(BUILD_UNIT_TESTS) enable_testing() 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_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_STORE_C_SHARED + "build a self-contained libmooncake_store.so exposing only the store_c.h C ABI (for dlopen consumers)" + OFF) option(WITH_EP "build mooncake with expert parallelism support" OFF) +option(USE_NOF "build mooncake store with NoF SSD pool support" OFF) +option(MOONCAKE_ENABLE_TEST_FAILPOINTS + "Enable file-handshake failpoints for integration tests" OFF) +if(MOONCAKE_ENABLE_TEST_FAILPOINTS) + add_compile_definitions(MOONCAKE_ENABLE_TEST_FAILPOINTS) +endif() +option(MOONCAKE_ENABLE_OPLOG_PERF_METRICS + "Enable detailed batch OpLog performance metrics" OFF) -add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extern/pybind11) -set(PYTHON_EXECUTABLE "python3") +include(${CMAKE_CURRENT_SOURCE_DIR}/mooncake-common/SetupPython.cmake) +if(SKBUILD) + set(PYBIND11_FINDPYTHON ON) + find_package(pybind11 CONFIG REQUIRED) +else() + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extern/pybind11) +endif() execute_process( - COMMAND ${PYTHON_EXECUTABLE} -c "import sys; print(sys.path[-1])" - OUTPUT_VARIABLE PYTHON_SYS_PATH -) + COMMAND ${PYTHON_EXECUTABLE} -c "import sys; print(sys.path[-1])" + OUTPUT_VARIABLE PYTHON_SYS_PATH) string(STRIP ${PYTHON_SYS_PATH} PYTHON_SYS_PATH) -if (USE_ETCD) +if(USE_ETCD) add_compile_definitions(USE_ETCD) - if (USE_ETCD_LEGACY) + if(USE_ETCD_LEGACY) add_compile_definitions(USE_ETCD_LEGACY) - message(STATUS "etcd as metadata server support is enabled (etcd-cpp-api-v3)") + message( + STATUS "etcd as metadata server support is enabled (etcd-cpp-api-v3)") else() message(STATUS "etcd as metadata server support is enabled (go package)") endif() endif() option(STORE_USE_ETCD "build mooncake store with etcd" OFF) -if (STORE_USE_ETCD) +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(ENABLE_KV_EVENTS + "Build master KV events ZMQ publisher (requires libzmq when ON)" OFF) +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() + +if(USE_NOF) + add_compile_definitions(USE_NOF) +else() + message(STATUS "USE_NOF=OFF, NoF SSD pool support is disabled") +endif() option(STORE_USE_JEMALLOC "Use jemalloc in mooncake store master" OFF) -# Define ASIO macros before adding mooncake-asio subdirectory +include(${CMAKE_CURRENT_SOURCE_DIR}/cmake/Dependencies.cmake) +if(WITH_STORE) + mooncake_provide_zstd() + mooncake_provide_xxhash() + if(ENABLE_KV_EVENTS) + mooncake_provide_libzmq() + endif() +endif() +if(WITH_STORE OR (WITH_TE AND USE_TENT)) + mooncake_provide_liburing() +endif() +if((WITH_TE AND USE_REDIS) OR (WITH_STORE AND STORE_USE_REDIS)) + mooncake_provide_hiredis(REQUIRED) +elseif(WITH_TE AND USE_TENT) + mooncake_provide_hiredis() +endif() + +# Define ASIO macros before building targets that include ASIO headers. add_compile_definitions(ASIO_SEPARATE_COMPILATION ASIO_DYN_LINK) -add_subdirectory(mooncake-asio) add_subdirectory(mooncake-common) include_directories(mooncake-common/etcd) +include_directories(mooncake-common/k8s-lease) include_directories(mooncake-common/include) -if (WITH_TE) +if(WITH_TE) add_subdirectory(mooncake-transfer-engine) include_directories(mooncake-transfer-engine/include) endif() -if (WITH_STORE) +if(WITH_STORE) message(STATUS "Mooncake Store will be built") add_subdirectory(mooncake-store) include_directories(mooncake-store/include) endif() -if (WITH_EP) - 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}") +if(WITH_STORE_RUST) + if(NOT WITH_STORE) + message(FATAL_ERROR "WITH_STORE_RUST=ON requires WITH_STORE=ON") endif() - set(EP_TORCH_VERSIONS "${EP_TORCH_VERSIONS}" CACHE STRING - "PyTorch versions for EP/PG extensions, semicolon-separated (empty = use currently-installed torch)") + message(STATUS "Mooncake Store Rust bindings will be built") + add_subdirectory(mooncake-store/rust) +endif() - # 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") +option(EP_USE_IDE "Enable intelligent indexing for IDEs" OFF) +if(WITH_EP) + add_subdirectory(mooncake-pg) + include_directories(mooncake-pg/include) + + 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_library( + mooncake_pg_torch_ide OBJECT + mooncake-pg/torch/src/mooncake_backend.cpp + mooncake-pg/torch/src/pg_py.cpp mooncake-pg/torch/src/work_handles.cpp) + # Reuse Mooncake EP's dependency include paths for clangd indexing. + target_include_directories( + mooncake_pg_torch_ide + PRIVATE mooncake-pg/torch/include mooncake-pg/include + $) + else() + message( + STATUS + "WITH_EP enabled: building Mooncake EP natively and PG via setup.py") + if(USE_CUDA) + find_package(CUDAToolkit REQUIRED) + message(STATUS "Detected CUDA version: ${CUDAToolkit_VERSION}") + endif() + + add_subdirectory(mooncake-ep) + include_directories(mooncake-ep/include) + + # 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)" + ) + + if(SKBUILD AND NOT EP_TORCH_VERSIONS) + execute_process( + COMMAND ${Python3_EXECUTABLE} -c "import torch" + RESULT_VARIABLE _mooncake_python_has_torch + OUTPUT_QUIET ERROR_QUIET) + if(NOT _mooncake_python_has_torch EQUAL 0) + message( + FATAL_ERROR + "WITH_EP requires EP_TORCH_VERSIONS in an isolated scikit-build-core build because PyTorch is not installed in the build environment" + ) + endif() + endif() + + # 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) + if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL "13.0") + set(TORCH_CUDA_ARCH_LIST "8.0;9.0;10.3") + else() + set(TORCH_CUDA_ARCH_LIST "8.0;9.0") + endif() + 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\")") + + # Keep generated Torch extensions and device libraries in backend-owned + # build directories. The legacy release path injects the staged device + # artifacts after auditwheel; scikit-build-core installs them directly into + # its private wheel tree. + set(EP_PG_STAGING_DIR "${CMAKE_BINARY_DIR}/ep_pg_staging") + set(EP_PG_BUILD_DIR "${CMAKE_BINARY_DIR}/ep_pg_build") + file(REMOVE_RECURSE "${EP_PG_STAGING_DIR}" "${EP_PG_BUILD_DIR}") + file(MAKE_DIRECTORY "${EP_PG_STAGING_DIR}" "${EP_PG_BUILD_DIR}") + + # 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}") + + if(USE_CUDA) + add_custom_target( + mooncake_ep_device_stage ALL + COMMAND ${CMAKE_COMMAND} -E make_directory "${EP_PG_STAGING_DIR}" + COMMAND ${CMAKE_COMMAND} -E copy_if_different + "$" "${EP_PG_STAGING_DIR}/" + COMMENT "Staging Mooncake EP CUDA device library" + DEPENDS mooncake_ep_device + VERBATIM) + endif() + + 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/torch" + "-DEP_CUDA_MAJOR=${CUDAToolkit_VERSION_MAJOR}" + "-DEP_CUDA_MINOR=${CUDAToolkit_VERSION_MINOR}" + "-DEP_TORCH_VERSIONS=${_ep_torch_versions_pipe}" + "-DSTAGING_DIR=${EP_PG_STAGING_DIR}" "-DBUILD_DIR=${EP_PG_BUILD_DIR}" + "-DPG_CORE_SO_PATH=$" + "-DPG_DEVICE_SO_PATH=$" + "-DPython3_EXECUTABLE=${Python3_EXECUTABLE}" + "-DEP_USE_MUSA=$,1,0>" + "-DEP_USE_MACA=$,1,0>" -P + "${CMAKE_CURRENT_SOURCE_DIR}/mooncake-pg/torch/BuildPgExt.cmake" + COMMENT "Building Mooncake PG Python extension(s)" + DEPENDS mooncake_pg mooncake_pg_device + VERBATIM) 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_TORCH_VERSIONS=${_ep_torch_versions_pipe}" - "-DTORCH_CUDA_ARCH_LIST=${_torch_cuda_arch_list_pipe}" - "-DSTAGING_DIR=${EP_PG_STAGING_DIR}" - "-DENGINE_SO_PATH=$" - -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_TORCH_VERSIONS=${_ep_torch_versions_pipe}" - "-DTORCH_CUDA_ARCH_LIST=${_torch_cuda_arch_list_pipe}" - "-DSTAGING_DIR=${EP_PG_STAGING_DIR}" - "-DENGINE_SO_PATH=$" - -P "${CMAKE_CURRENT_SOURCE_DIR}/mooncake-pg/BuildPgExt.cmake" - COMMENT "Building Mooncake PG Python extension(s)" - DEPENDS engine mooncake_ep_ext - VERBATIM - ) endif() add_subdirectory(mooncake-integration) -if (WITH_P2P_STORE) +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() diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 54d88c2a90..fae47e6fed 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -16,7 +16,8 @@ Thank you for your interest in contributing to Mooncake! Our community warmly we ### PR Title and Classification -Use a prefixed PR title to indicate the type of changes. Please use one of the following: +Use a prefixed PR title to indicate the type or module affected by the changes. +Prefer one of the following documented prefixes: - ``[Bugfix]`` for bug fixes. - ``[CI/Build]`` for build or continuous integration improvements. @@ -28,6 +29,11 @@ Use a prefixed PR title to indicate the type of changes. Please use one of the f - ``[Misc]`` for PRs that do not fit the above categories. Please use this sparingly. +The project history also contains common aliases and module prefixes. Use these +when they better match the change scope: ``[Bug fix]``, ``[Build]``, ``[CI]``, +``[Docs]``, ``[EP]``, ``[Feature]``, ``[MUSA]``, ``[PG]``, ``[TE]``, +``[TENT]``, and ``[Wheel]``. + ### RFC Discussion For major architectural changes (>500 LOC excluding tests), we would expect a GitHub issue (RFC) discussing the technical design and justification. @@ -41,20 +47,39 @@ 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 --staged` | Format staged C/C++ changes 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` | | CMake | cmake-format | Keep build scripts readable | | Meta | check-yaml / check-merge-conflict / check-added-large-files | Prevent bad commits | #### Setup ```bash -pip install -r requirements-dev.txt +pip install -r requirements.txt pre-commit install ``` +After installation, every commit formats only the added or modified lines in +staged C/C++ files. If the hook rewrites a file, review and re-stage it before +committing again. Use `./scripts/code_format.sh --all` only when intentionally +formatting the whole project. + #### Usage -Run on all files (first run will install hook environments): +After `pre-commit install`, hooks run on each commit. To run them manually on +staged files (the first run may install hook environments): +```bash +pre-commit run +``` +Before opening a PR, run hooks only on files changed against the PR base +(default `origin/main`). The C/C++ hook remains limited to staged or changed +line ranges: +```bash +git fetch origin main +pre-commit run --files $(git diff --name-only --diff-filter=ACMR origin/main...HEAD) +``` +Use full-repo checks only for intentional whole-project cleanup; do not fold +unrelated rewrites into a feature PR. Prefer `./scripts/code_format.sh --all` +for a deliberate whole-project C/C++ format: ```bash pre-commit run --all-files ``` @@ -65,9 +90,10 @@ git add .pre-commit-config.yaml git commit -m "chore: pre-commit autoupdate" ``` -If clang-format is missing, install it (Ubuntu example): +If clang-format or its `git-clang-format` helper is missing, install the LLVM +20 package (Ubuntu example): ```bash -sudo apt-get update && sudo apt-get install -y clang-format +sudo apt-get update && sudo apt-get install -y clang-format-20 ``` You can temporarily skip hooks: @@ -77,7 +103,15 @@ git commit -m "wip: skipping hooks" --no-verify But please avoid using `--no-verify` for routine commits to keep code quality high. #### CI Integration -The configuration supports automatic fixing PRs via `pre-commit.ci` if enabled. To activate, add the repository in the pre-commit.ci dashboard; no further changes are needed. +GitHub pull-request and push checks validate only added or modified C/C++ line +ranges relative to the selected base revision. This avoids failing a focused +change solely because an otherwise untouched part of the same file has older +formatting. Changed-line selection is delegated to LLVM's `git-clang-format` +helper so the local hook and CI use the same Git-aware behavior. + +The configuration also supports automatic fixing PRs via `pre-commit.ci` if +enabled. To activate, add the repository in the pre-commit.ci dashboard; no +further changes are needed. ## Code Quality diff --git a/FAST25-release/README.md b/FAST25-release/README.md new file mode 100644 index 0000000000..c4460e62a8 --- /dev/null +++ b/FAST25-release/README.md @@ -0,0 +1,59 @@ +# Mooncake FAST'25 Trace Release + +This directory contains the open-source request traces associated with Mooncake. + +The `traces/` directory is the updated trace release used by the FAST'25 paper. The older `arxiv-trace/mooncake_trace.jsonl` file is the historical single-file trace release from the arXiv technical report. + +## Directory Layout + +| Path | Description | +| --- | --- | +| `traces/conversation_trace.jsonl` | FAST'25 conversation workload trace. | +| `traces/toolagent_trace.jsonl` | FAST'25 tool and agent workload trace. | +| `traces/synthetic_trace.jsonl` | FAST'25 synthetic workload trace built from public datasets. | +| `arxiv-trace/mooncake_trace.jsonl` | Historical trace released with the arXiv technical report. | +| `Mooncake-FAST25.pdf` | FAST'25 paper with the trace appendix. | + +For new experiments that reproduce or build on the FAST'25 paper, prefer the three files under `traces/`. Use `arxiv-trace/mooncake_trace.jsonl` when referring to the original arXiv technical report trace or comparing against the earlier single-file trace release. + +## Workloads + +The FAST'25 trace release contains three workloads: + +| Workload | File | Requests | Avg input length | Avg output length | Arrival pattern | +| --- | --- | ---: | ---: | ---: | --- | +| Conversation | `traces/conversation_trace.jsonl` | 12,031 | 12,035 | 343 | Timestamp | +| Tool and Agent | `traces/toolagent_trace.jsonl` | 23,608 | 8,596 | 182 | Timestamp | +| Synthetic | `traces/synthetic_trace.jsonl` | 3,993 | 15,325 | 149 | Poisson | + +The conversation and tool and agent traces are sampled from one hour of online request data. The synthetic trace is constructed from public datasets and uses Poisson-generated arrivals while preserving the order within multi-turn conversations. + +## Trace Format + +Each trace is a JSONL file. Each line is one request: + +```json +{ + "timestamp": 27482, + "input_length": 6955, + "output_length": 52, + "hash_ids": [46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 2353, 2354] +} +{ + "timestamp": 30535, + "input_length": 6472, + "output_length": 26, + "hash_ids": [46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 2366] +} +``` + +Fields: + +- `timestamp`: Relative request arrival time in milliseconds. In the real traces, requests should be replayed according to these timestamps. In the synthetic trace, timestamps are generated from a Poisson process. +- `input_length`: Number of input tokens. Raw text and tokens are not included for privacy protection. +- `output_length`: Number of output tokens. +- `hash_ids`: Remapped prefix block hashes. The block size is 512 tokens. Identical hash IDs indicate reusable prefix KV cache blocks. For example, the two sample requests share the first 12 hash IDs, so they can share prefix caching for the first `12 * 512 = 6144` tokens. + +The trace files contain only anonymized timing, length, and remapped hash information. They are intended for reproducible simulation and evaluation of KV cache reuse behavior without exposing user content. + +For more details, see the FAST'25 paper appendix in [`Mooncake-FAST25.pdf`](Mooncake-FAST25.pdf). diff --git a/MAINTAINERS.md b/MAINTAINERS.md index 3e978d9c9c..24260e376e 100644 --- a/MAINTAINERS.md +++ b/MAINTAINERS.md @@ -12,6 +12,7 @@ Current list of codeowners on this project: | ---------------------------- | ---------------------------- | ---------------------------- | ---------------------------- | | | | | | | | | | | -| | | | | +| | | | | +| | | | | Want to include your company logo? Just open a Pull Request! \ No newline at end of file diff --git a/README.md b/README.md index 3bafba7653..30d202f3f2 100644 --- a/README.md +++ b/README.md @@ -6,32 +6,44 @@ Paper | Slides | Traces - | Technical Report - | Blog + | Documentation + | Blog | Slack

- [![Docs](https://img.shields.io/badge/docs-live-brightgreen)](https://kvcache-ai.github.io/Mooncake/) - [![PyPI](https://img.shields.io/pypi/v/mooncake-transfer-engine)](https://pypi.org/project/mooncake-transfer-engine) - [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/mooncake-transfer-engine)](https://pypi.org/project/mooncake-transfer-engine) - [![CUDA <=12.9](https://img.shields.io/static/v1?label=CUDA&message=%3C%3D12.9&color=76B900)](https://pypi.org/project/mooncake-transfer-engine) - [![CUDA 13.0/13.1](https://img.shields.io/static/v1?label=CUDA&message=13.0%2F13.1&color=76B900)](https://pypi.org/project/mooncake-transfer-engine-cuda13) - [![PyPI - Downloads](https://img.shields.io/pypi/dm/mooncake-transfer-engine)](https://pypi.org/project/mooncake-transfer-engine) [![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/kvcache-ai/Mooncake) + [![PyPI - Downloads](https://static.pepy.tech/badge/mooncake-transfer-engine?period=month)](https://pypi.org/project/mooncake-transfer-engine) [![GitHub commit activity](https://img.shields.io/github/commit-activity/w/kvcache-ai/Mooncake)](https://github.com/kvcache-ai/Mooncake/graphs/commit-activity) [![license](https://img.shields.io/github/license/kvcache-ai/mooncake.svg)](https://github.com/kvcache-ai/Mooncake/blob/main/LICENSE-APACHE) + [![Docker](https://img.shields.io/docker/v/kvcacheai/mooncake?label=docker&logo=docker&logoColor=white&color=2496ED)](https://hub.docker.com/r/kvcacheai/mooncake) +
+ [![PyPI CUDA <=12.9](https://img.shields.io/static/v1?label=pypi&message=CUDA%20%3C%3D12.9&color=76B900)](https://pypi.org/project/mooncake-transfer-engine) + [![PyPI CUDA 13.0/13.1](https://img.shields.io/static/v1?label=pypi&message=CUDA%2013.0%2F13.1&color=76B900)](https://pypi.org/project/mooncake-transfer-engine-cuda13) + [![PyPI Non-CUDA](https://img.shields.io/static/v1?label=pypi&message=non-CUDA&color=00BFFF)](https://pypi.org/project/mooncake-transfer-engine-non-cuda/) + [![PyPI NPU](https://img.shields.io/static/v1?label=pypi&message=NPU&color=F87171)](https://pypi.org/project/mooncake-transfer-engine-npu/) + [![PyPI MUSA](https://img.shields.io/static/v1?label=pypi&message=MUSA&color=F97316)](https://pypi.org/project/mooncake-transfer-engine-musa/) + [![PyPI EFA](https://img.shields.io/static/v1?label=pypi&message=EFA&color=F59E0B)](https://pypi.org/project/mooncake-transfer-engine-efa/) + [![PyPI ROCm](https://img.shields.io/static/v1?label=pypi&message=ROCm&color=black)](https://pypi.org/project/mooncake-transfer-engine-rocm/)
Mooncake is the serving platform for icon Kimi, a leading LLM service provided by icon Moonshot AI. -Now both the Transfer Engine and Mooncake Store are open-sourced! -This repository also hosts its technical report and the open-sourced traces. +Under real workloads, Mooncake’s innovative architecture enables Kimi to handle 75% more requests while adhering to SLOs.

🔄 Updates

-- **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. +- **Aug 20, 2026**: Mooncake is integrated into [Miles](https://github.com/radixark/miles) as a rollout data-transfer backend for the fragmented, heterogeneous data moving between rollout and training in disaggregated RL. Blogs: [KVCache.AI](https://kvcache.ai/blog/mooncake-rl-rollout-data-transfer/), [lmsys](https://www.lmsys.org/blog/2026-08-20-miles-mooncake-rollout-data-transfer). +- **Aug 17, 2026**: Mooncake is integrated into [Speculators](https://github.com/vllm-project/speculators) as a distributed backend for multi-node online training, efficiently moves hidden-state between vLLM inference workers and Speculators trainers through RDMA, eliminating the need for massive hidden-state storage in offline training. [Benchmark on GB300 NVL72](https://x.com/mgoin_/status/2072785822231728363). +- **May 7, 2026**: 🚀 [vLLM officially features Mooncake Store](https://vllm.ai/blog/mooncake-store) — a deep dive into how Mooncake's distributed KVCache engine supercharges vLLM inference with high-throughput, memory-efficient, cross-instance KV cache sharing! +- **Apr 29, 2026**: SGLang introduces [RDMA-based P2P weight transfer for large-scale distributed RL](https://lmsys.org/blog/2026-04-29-p2p-update/) using Mooncake TransferEngine, achieving 7x faster weight updates for the 1T-parameter Kimi-K2 model (53s → 7.2s) with zero-copy RDMA transfer across thousands of GPUs. +- **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. + +
+More + +- **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. Details in [blog](https://light-ai.top/LightX2V-BLOG/posts/Disaggregation/). - **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! @@ -47,9 +59,9 @@ This repository also hosts its technical report and the open-sourced traces. - **Aug 23, 2025**: [xLLM](https://github.com/jd-opensource/xllm) high-performance inference engine builds hybrid KV cache management based on Mooncake, supporting global KV cache management with intelligent offloading and prefetching. - **Aug 18, 2025**: vLLM-Ascend [integrates Mooncake Transfer Engine](https://docs.vllm.ai/projects/ascend/en/latest/developer_guide/feature_guide/disaggregated_prefill.html) for KV cache register and disaggregate prefill, enabling efficient distributed inference on Ascend NPUs. - **Jul 20, 2025**: Mooncake powers [the deployment of Kimi K2](https://lmsys.org/blog/2025-07-20-k2-large-scale-ep/) on 128 H200 GPUs with PD disaggregation and large-scale expert parallelism, achieving 224k tokens/sec prefill throughput and 288k tokens/sec decode throughput. - - **Jun 20, 2025**: Mooncake becomes a PD disaggregation [backend](https://kvcache-ai.github.io/Mooncake/getting_started/examples/lmdeploy-integration-v0.9.html) for LMDeploy. + - **Jun 20, 2025**: Mooncake becomes a PD disaggregation [backend](https://kvcache-ai.github.io/Mooncake/deployment/integrations/lmdeploy.html) for LMDeploy. - **May 9, 2025**: NIXL officially supports Mooncake Transfer Engine as [a backend plugin](https://github.com/ai-dynamo/nixl/blob/main/src/plugins/mooncake/README.md). - - **May 8, 2025**: [Mooncake x LMCache](https://kvcache-ai.github.io/Mooncake/getting_started/examples/lmcache-integration.html) unite to pioneer KVCache-centric LLM serving system. + - **May 8, 2025**: [Mooncake x LMCache](https://kvcache-ai.github.io/Mooncake/deployment/integrations/lmcache/index.html) unite to pioneer KVCache-centric LLM serving system. - **May 5, 2025**: Supported by Mooncake Team, SGLang release guidance to deploy DeepSeek with PD Disaggregation on 96 H100 GPUs. - **Apr 22, 2025**: LMCache officially supports Mooncake Store as a remote connector. - **Apr 10, 2025**: SGLang officially supports Mooncake Transfer Engine for disaggregated prefilling and KV cache transfer. @@ -62,126 +74,136 @@ This repository also hosts its technical report and the open-sourced traces. - **June 27, 2024**: We present a series of Chinese blogs with more discussions on zhihu 1, 2, 3, 4, 5, 6, 7. - **June 26, 2024**: Initial technical report release. +

🎉 Overview

-Mooncake features a KVCache-centric disaggregated architecture that separates the prefill and decoding clusters. It also leverages the underutilized CPU, DRAM, and SSD resources of the GPU cluster to implement a disaggregated KVCache pool. + +
+ +
-![architecture](image/architecture.png) +Mooncake is an infrastructure project for large-scale LLM inference and training. It features a KV cache-centric disaggregated architecture that separates prefill and decode clusters, while leveraging otherwise underutilized CPU, DRAM, and SSD resources in GPU clusters to build a disaggregated KV cache pool. -The core of Mooncake is its KVCache-centric scheduler, which balances maximizing overall effective throughput while meeting latency-related Service Level Objectives (SLOs). Unlike traditional studies that assume all requests will be processed, Mooncake faces challenges in highly overloaded scenarios. To mitigate these, we developed a prediction-based early rejection policy. Experiments show that Mooncake excels in long-context scenarios. Compared to the baseline method, Mooncake can achieve up to a 525% increase in throughput in certain simulated scenarios while adhering to SLOs. Under real workloads, Mooncake’s innovative architecture enables Kimi to handle 75% more requests. +Mooncake includes a high-performance Transfer Engine for low-latency data movement across heterogeneous networks and accelerators; Mooncake Store for distributed KV cache and model-weight management; and Mooncake EP & PG for elastic MoE serving. Deeply integrated with ecosystems such as SGLang and vLLM, Mooncake helps LLM systems improve cache reuse, reduce serving latency, and scale efficiently across multi-node clusters. -

🧩 Components

+

🔥 Show Cases

- - +### Transfer Engine (TE) -**Mooncake Core Component: Transfer Engine (TE)** -The core of Mooncake is the Transfer Engine (TE), which provides a unified interface for batched data transfer across various storage devices and network links. Supporting multiple protocols including TCP, RDMA, CXL/shared-memory, and NVMe over Fabric (NVMe-of), TE is designed to enable fast and reliable data transfer for AI workloads. Compared to Gloo (used by Distributed PyTorch) and traditional TCP, TE achieves significantly lower I/O latency, making it a superior solution for efficient data transmission. +The core of Mooncake is the Transfer Engine (TE), a high-performance data transfer framework. TE offers a unified interface for batched data movement across diverse storage, network, and accelerator environments. By supporting multiple transport protocols, topology-aware routing, multi-NIC bandwidth aggregation, and automatic failover, TE delivers low-latency, scalable, and robust data transmission for distributed AI workloads. See the [Transfer Engine guide](https://kvcache-ai.github.io/Mooncake/design/transfer-engine/index.html) for details. -**P2P Store and Mooncake Store** -Both P2P Store and Mooncake Store are built on the Transfer Engine and provide key/value caching for different scenarios. P2P Store focuses on sharing temporary objects (e.g., checkpoint files) across nodes in a cluster, preventing bandwidth saturation on a single machine. Mooncake Store, on the other hand, supports distributed pooled KVCache, specifically designed for XpYd disaggregation to enhance resource utilization and system performance. +
+Highlights -**Mooncake Integration with Leading LLM Inference Systems** -Mooncake has been seamlessly integrated with several popular large language model (LLM) inference systems. Through collaboration with the vLLM and SGLang teams, Mooncake now officially supports prefill-decode disaggregation. By leveraging the high-efficiency communication capabilities of RDMA devices, Mooncake significantly improves inference efficiency in prefill-decode disaggregation scenarios, providing robust technical support for large-scale distributed inference tasks. -In addition, Mooncake has been successfully integrated with SGLang's Hierarchical KV Caching, vLLM's prefill serving, and LMCache, augmenting KV cache management capabilities across large-scale inference scenarios. +- **Efficient use of multiple RDMA NIC devices.** Transfer Engine supports the use of multiple RDMA NIC devices to achieve the *aggregation of transfer bandwidth*. -**Elastic Expert Parallelism Support** -Mooncake adds elasticity and fault tolerance support for MoE model inference, enabling inference systems to remain responsive and recoverable in the event of GPU failures or changes in resource configuration. This functionality includes automatic faulty rank detection and can work with the EPLB module to dynamically route tokens to healthy ranks during inference. +- **Topology-aware path selection.** Transfer Engine can *select optimal devices* based on the location (NUMA affinity, etc.) of both source and destination. -**Tensor-Centric Ecosystem** -Mooncake establishes a full-stack, Tensor-oriented AI infrastructure where Tensors serve as the fundamental data carrier. The ecosystem spans from the Transfer Engine, which accelerates Tensor data movement across heterogeneous storage (DRAM/VRAM/NVMe), to the P2P Store and Mooncake Store for distributed management of Tensor objects (e.g., Checkpoints and KVCache), up to the Mooncake Backend enabling Tensor-based elastic distributed computing. This architecture is designed to maximize Tensor processing efficiency for large-scale model inference and training. +- **Robust against temporary network errors.** Once transmission fails, Transfer Engine will try to use alternative paths for data delivery automatically. -

🔥 Show Cases

+- **Superior performance at scale.** With 40 GB of data (equivalent to the size of the KVCache generated by 128k tokens in the LLaMA3-70B model), Mooncake Transfer Engine delivers up to **87 GB/s** and **190 GB/s** of bandwidth in 4×200 Gbps and 8×400 Gbps RoCE networks respectively, which are about **2.4x and 4.6x faster** than the TCP protocol. -### 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. 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). +- **Broad support for heterogeneous transports and accelerators.** Transfer Engine provides unified data transfer across diverse protocols, including TCP, RDMA, AWS EFA, NVMe-oF, NVLink, HIP, Barex, CXL, and Ascend-family transports. When built with the corresponding runtime, Transfer Engine can detect accelerator memory and select suitable transport paths for efficient data movement across CUDA, MUSA, HIP, MACA, Cambricon MLU, and Ascend-enabled environments. 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*. +- **Widely adopted across the LLM ecosystem.** TE is used in production inference stacks such as [SGLang](https://github.com/sgl-project/sglang), [vLLM](https://github.com/vllm-project/vllm), [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM), [vLLM-Ascend](https://github.com/vllm-project/vllm-ascend), [checkpoint-engine](https://github.com/MoonshotAI/checkpoint-engine), and [NIXL](https://github.com/ai-dynamo/nixl), among others, to efficiently transfer KV cache, embeddings, model weights, and other data. -- **Topology aware path selection.** Transfer Engine can *select optimal devices* based on the location (NUMA affinity, etc.) of both source and destination. +
-- **More robust against temporary network errors.** Once transmission fails, Transfer Engine will try to use alternative paths for data delivery automatically. +### Mooncake Store -#### Performance -With 40 GB of data (equivalent to the size of the KVCache generated by 128k tokens in the LLaMA3-70B model), Mooncake Transfer Engine delivers up to **87 GB/s** and **190 GB/s** of bandwidth in 4×200 Gbps and 8×400 Gbps RoCE networks respectively, which are about **2.4x and 4.6x faster** than the TCP protocol. +Mooncake Store is a high-performance distributed key-value cache storage engine designed for LLM inference. Built on the Transfer Engine, it stores and manages reusable KV caches and model weights across inference clusters, with support for efficient object storage, replication, eviction, and high-bandwidth data transfer. See the [Mooncake Store guide](https://kvcache-ai.github.io/Mooncake/design/mooncake-store.html) for details. - - +
+Highlights + +- **High bandwidth utilization.** Mooncake Store supports large-object striping, parallel I/O, and end-to-end zero-copy data transfer, fully utilizing aggregated bandwidth across multiple NICs. + +- **Multi-tier cache hierarchy**. Mooncake Store supports a multi-level cache design across DRAM and SSD/NVMe, enabling larger cache capacity. + +- **Elastic and disaggregated storage.** Mooncake Store decouples KVCache storage from inference engines, allowing storage nodes to be dynamically added or removed while keeping cached data independent from engine restarts, upgrades, and scheduling decisions. + +- **Programmatic object management.** Mooncake Store allows applications to control object placement and lifecycle through per-object policies, including replica counts, preferred segments, soft pin, and hard pin. These controls help inference systems protect important KV caches and model weights while guiding replication, placement, and eviction behavior. + +- **Broad ecosystem adoption.** Mooncake Store is used across the LLM systems ecosystem as a high-performance distributed storage backend for KV caches, hidden states, and model weights. It supports integrations with [SGLang's Hierarchical KV Caching](https://lmsys.org/blog/2025-09-10-sglang-hicache/), [vLLM's prefill serving](https://docs.vllm.ai/en/latest/features/disagg_prefill.html), and [LMCache](https://kvcache-ai.github.io/Mooncake/deployment/integrations/lmcache/index.html), and has been adopted by systems such as [TorchSpec](https://pytorch.org/blog/torchspec-speculative-decoding-training-at-scale/) and [TransferQueue](https://github.com/Ascend/TransferQueue) to decouple inference, training, and reinforcement-learning workloads through efficient state management and asynchronous data movement. -### P2P Store ([Guide](https://kvcache-ai.github.io/Mooncake/design/p2p-store.html)) -P2P Store is built on the Transfer Engine and supports sharing temporary objects between peer nodes in a cluster. P2P Store is ideal for scenarios like checkpoint transfer, where data needs to be rapidly and efficiently shared across a cluster. -**P2P Store has been used in the checkpoint transfer service of Moonshot AI.** +
-#### Highlights -- **Decentralized architecture.** P2P Store leverages a pure client-side architecture with global metadata managed by the etcd service. +### Mooncake EP and Process Group (PG) -- **Efficient data distribution.** Designed to enhance the efficiency of large-scale data distribution, P2P Store *avoids bandwidth saturation* issues by allowing replicated nodes to share data directly. This reduces the CPU/RDMA NIC pressures of data providers (e.g., trainers). +Mooncake EP and Mooncake PG extend Mooncake from high-performance data movement to fault-tolerant distributed execution for large-scale MoE inference. Mooncake EP adapts DeepEP-style expert-parallel dispatch and combine operations with rank activeness awareness, while Mooncake PG provides a PyTorch distributed process-group backend with collective communication primitives that can detect failed ranks, report failures to upper layers, and recover ranks without restarting the entire inference service. See the [Mooncake EP & Backend guide](https://kvcache-ai.github.io/Mooncake/api-reference/python/ep-backend.html) for details. - +
+Highlights - +- **Fault-tolerant expert parallelism.** Mooncake EP adds `active_ranks` awareness to expert-parallel dispatch and combine APIs, allowing MoE inference systems to route around failed ranks and continue serving with healthy experts. -### Mooncake Store ([Guide](https://kvcache-ai.github.io/Mooncake/design/mooncake-store.html)) -Mooncake Store is a distributed KVCache storage engine specialized for LLM inference based on Transfer Engine. It is the central component of the KVCache-centric disaggregated architecture. The goal of Mooncake Store is to store the reusable KV caches across various locations in an inference cluster. Mooncake Store has been supported in [SGLang's Hierarchical KV Caching](https://lmsys.org/blog/2025-09-10-sglang-hicache/), [vLLM's prefill serving](https://docs.vllm.ai/en/latest/features/disagg_prefill.html) and is now integrated with [LMCache](https://kvcache-ai.github.io/Mooncake/getting_started/examples/lmcache-integration.html) to provide enhanced KVCache management capabilities. +- **DeepEP-compatible programming model.** Mooncake EP keeps the API largely consistent with DeepEP's low-latency mode, making it easier for inference engines to adopt fault-tolerant expert parallelism without rewriting their MoE communication stack. -#### Highlights -- **Multi-replica support**: Mooncake Store supports storing multiple data replicas for the same object, effectively alleviating hotspots in access pressure. +- **PyTorch ProcessGroup integration.** Mooncake PG can be registered as a `torch.distributed` backend, enabling standard collective APIs such as `all_gather` while using Mooncake's communication and failure-reporting mechanisms underneath. -- **High bandwidth utilization**: Mooncake Store supports striping and parallel I/O transfer of large objects, fully utilizing multi-NIC aggregated bandwidth for high-speed data reads and writes. +- **Elastic rank recovery.** Mooncake PG exposes recovery-oriented primitives such as peer-state polling and rank recovery, allowing replacement processes to rejoin existing process groups and helping inference services recover from partial failures. -### SGLang Integration ([Guide](https://kvcache-ai.github.io/Mooncake/getting_started/examples/sglang-integration/hicache-integration-v1.html)) +- **SGLang integration for production MoE serving.** Mooncake's collective backend and expert-parallel kernels are integrated into SGLang to support fault-tolerant expert-parallel inference for large MoE models, including Elastic Expert Parallel serving scenarios. -SGLang officially supports Mooncake Store as a [HiCache storage backend](https://lmsys.org/blog/2025-09-10-sglang-hicache/). This integration enables scalable KV cache retention and high-performance access for large-scale LLM serving scenarios. +
+ +### Tensor-Centric Ecosystem + +Mooncake establishes a full-stack, Tensor-oriented AI infrastructure where Tensors serve as the fundamental data carrier. The ecosystem spans from the Transfer Engine, which accelerates Tensor data movement across heterogeneous storage (DRAM/VRAM/NVMe), to Mooncake Store for distributed management of Tensor objects (e.g., KVCache and model weight), up to the Mooncake Backend enabling Tensor-based elastic distributed computing. This architecture is designed to maximize Tensor processing efficiency for large-scale model inference and training. + +### SGLang Integration ([Guide](https://kvcache-ai.github.io/Mooncake/deployment/integrations/sglang/index.html)) + +Mooncake is deeply integrated into [SGLang](https://github.com/sgl-project/sglang/) as a high-performance communication and storage backend. These integrations enable efficient KV cache transfer in PD-disaggregated serving, scalable multi-level KV caching through HiCache, fault-tolerant expert-parallel inference, high-performance multimodal pipeline data movement, and fast RDMA-based weight synchronization for large-scale RL training. Together, Mooncake and SGLang provide a production-oriented foundation for building elastic, high-throughput, and resource-efficient LLM and multimodal serving systems. + +
+Details + +- **PD Disaggregated Serving:** SGLang officially supports Mooncake Transfer Engine as a backend for disaggregated serving and KV cache transfer, enabling prefill and decode workers to exchange KV cache data efficiently across devices and machines. -#### Highlights - **Hierarchical KV Caching**: Mooncake Store serves as an external storage backend in SGLang's HiCache system, extending RadixAttention with multi-level KV cache storage across device, host, and remote storage layers. -- **Flexible Cache Management**: Supports multiple cache policies including write-through, write-through-selective, and write-back modes, with intelligent prefetching strategies for optimal performance. -- **Comprehensive Optimizations**: Features advanced data plane optimizations including page-first memory layout for improved I/O efficiency, zero-copy mechanisms for reduced memory overhead, GPU-assisted I/O kernels delivering fast CPU-GPU transfers, and layer-wise overlapping for concurrent KV cache loading while computation executes. -- **Elastic Expert Parallel**: Mooncake's collective communication backend and expert parallel kernels are integrated into SGLang to enable fault-tolerant expert parallel inference ([sglang#11657](https://github.com/sgl-project/sglang/pull/11657)). -- **Significant Performance Gains**: The multi-turn benchmark demonstrates substantial performance improvements over the non-HiCache setting. See our [benchmark report](https://kvcache-ai.github.io/Mooncake/performance/sglang-hicache-benchmark-results-v1.html) for more details. -- **Community Feedback**: Effective KV caching significantly reduces TTFT by eliminating redundant and costly re-computation. Integrating SGLang HiCache with the Mooncake service enables scalable KV cache retention and high-performance access. In our evaluation, we tested the DeepSeek-R1-671B model under PD-disaggregated deployment using in-house online requests sampled from a general QA scenario. On average, cache hits achieved an 84% reduction in TTFT compared to full re-computation. – Ant Group -### vLLM Integration ([Guide v0.2](https://kvcache-ai.github.io/Mooncake/getting_started/examples/vllm-integration/vllm-integration-v0.2.html)) -To optimize LLM inference, the vLLM community is working on supporting [disaggregated prefilling (PR 10502)](https://github.com/vllm-project/vllm/pull/10502). This feature allows separating the **prefill** phase from the **decode** phase in different processes. The vLLM uses `nccl` and `gloo` as the transport layer by default, but currently it cannot efficiently decouple both phases in different machines. +- **Elastic Expert Parallel**: Mooncake's collective communication backend and expert parallel kernels are integrated into SGLang to enable fault-tolerant expert parallel inference ([Elastic EP](https://www.lmsys.org/blog/2026-03-25-eep-partial-failure-tolerance/)). + +- **Cloud-Native SGLang HiCache Deployment with RBG**: The [RBG](https://github.com/sgl-project/rbg) + SGLang HiCache + Mooncake integration provides a role-based, out-of-the-box cloud-native deployment solution that is elastic, scalable, and optimized for high-performance inference workloads. + +- **Encode-Prefill-Decode Disaggregation for Multimodal Serving**: SGLang introduces Encode-Prefill-Decode disaggregation with Mooncake as a transfer backend. This enables compute-intensive multimodal encoders, such as Vision Transformers, to be decoupled from language model workers while transferring large embeddings efficiently through Mooncake’s RDMA-based engine. + +- **SGLang-Omni Multi-Stage Pipeline Data Transfer**: [SGLang-Omni](https://github.com/sgl-project/sglang-omni) integrates Mooncake as a relay backend for efficient cross-stage tensor and blob transfer in multimodal serving pipelines. This enables high-performance data movement between heterogeneous components such as thinker, talker, codec, and vocoder stages. + +- **RDMA-Based P2P Weight Transfer for Distributed RL**: SGLang adopts Mooncake TransferEngine for RDMA-based peer-to-peer weight transfer in large-scale distributed reinforcement learning. This enables zero-copy weight updates across thousands of GPUs and significantly accelerates synchronization for trillion-parameter models. -We have implemented vLLM integration, which uses Transfer Engine as the network layer instead of `nccl` and `gloo`, to support **inter-node KVCache transfer** [(PR 10884)](https://github.com/vllm-project/vllm/pull/10884). Transfer Engine provides simpler interfaces and more efficient use of RDMA devices. +
-We will soon release the new vLLM integration based on Mooncake Store, which supports xPyD prefill/decode disaggregation. +### vLLM Integration ([Guide](https://kvcache-ai.github.io/Mooncake/deployment/integrations/vllm/index.html)) -**_Update[Dec 16, 2024]: Here is the latest vLLM Integration ([Guide v0.2](https://kvcache-ai.github.io/Mooncake/getting_started/examples/vllm-integration/vllm-integration-v0.2.html)) that is based on vLLM's main branch._** +Mooncake integrates with [vLLM](https://github.com/vllm-project/vllm) to accelerate large language model serving through high-performance KV cache transfer and distributed KV cache storage. The integration supports both disaggregated prefill-decode serving and cross-instance KV cache sharing, helping vLLM deployments reduce TTFT, improve cache reuse, and scale more efficiently across multi-node inference clusters. -#### Performance -By supporting Topology Aware Path Selection and multi-card bandwidth aggregation, Mean TTFT of vLLM with Transfer Engine is up to 25% lower than traditional TCP-based transports. -In the future, we will further improve TTFT through GPUDirect RDMA and zero-copy. +
+Details -| Backend/Setting | Output Token Throughput (tok/s) | Total Token Throughput (tok/s) | Mean TTFT (ms) | Median TTFT (ms) | P99 TTFT (ms)| -|---------------------------------------------------------|---------------------------------|--------------------------------|----------------|------------------|---------------| -| Transfer Engine (RDMA) | 12.06 | 2042.74 | 1056.76 | 635.00 | 4006.59 | -| TCP | 12.05 | 2041.13 | 1414.05 | 766.23 | 6035.36 | +- **Disaggregated prefill-decode serving**: Mooncake enables vLLM to split prefill and decode workloads across different nodes. Through MooncakeConnector, vLLM transfers KV cache blocks from prefill workers to decode workers using Mooncake’s high-performance transfer engine, allowing prefill and decode resources to scale independently while keeping cross-node KV transfer overhead low. -- Click [here](https://kvcache-ai.github.io/Mooncake/performance/vllm-benchmark-results-v0.2.html) to access detailed benchmark results. +- **Distributed KV cache pooling and sharing**: [Mooncake Store extends vLLM](https://vllm.ai/blog/2026-05-06-mooncake-store) from isolated per-instance KV caches to a shared, cluster-level KV cache pool. Through MooncakeStoreConnector, multiple vLLM instances can store, retrieve, and reuse KV cache blocks based on hash-based prefix caching, reducing redundant prefill computation and improving cache efficiency for workloads with repeated prefixes, especially agentic and multi-turn serving scenarios. -**More advanced features are coming soon, so stay tuned!** +- **vLLM-Omni stage communication**: Mooncake also integrates with [vLLM-Omni](https://github.com/vllm-project/vllm-omni) through `MooncakeTransferEngineConnector` and `MooncakeStoreConnector`, enabling efficient cross-node data exchange between vLLM-Omni stages. -

🚀 Quick Start

+
-### Before using Mooncake +

🖥️ Supported Hardware

-Mooncake is designed and optimized for high-speed RDMA networks. Though Mooncake supports TCP-only data transfer, we **strongly** recommend users to evaluate the functionality and performance of Mooncake with RDMA network support. +Mooncake supports hardware backends across accelerator vendors, cloud fabrics, and standard datacenter interconnects, as listed below. See the [supported protocols](https://kvcache-ai.github.io/Mooncake/getting_started/supported-protocols.html) and [Transfer Engine design docs](https://kvcache-ai.github.io/Mooncake/design/transfer-engine/index.html) for details. -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)*. +| NVIDIA | Huawei | AMD | Cambricon | Moore Threads | AWS | +| --- | --- | --- | --- | --- | --- | +| MetaX | T-Head | Alibaba Cloud | Sunrise | Hygon | Biren Technology | -### Use Python package -The simplest way to use Mooncake Transfer Engine is using `pip`: +

🚀 Getting Started

-**For CUDA-enabled systems:** +Install Mooncake using `pip`. The `mooncake-transfer-engine` package includes Mooncake Transfer Engine, Mooncake Store, Mooncake EP and PG: - CUDA < 13.0 ```bash @@ -192,98 +214,84 @@ pip install mooncake-transfer-engine pip install mooncake-transfer-engine-cuda13 ``` -**For non-CUDA systems:** -```bash -pip install mooncake-transfer-engine-non-cuda -``` +In addition to CUDA, Mooncake also supports other accelerator backends, along with flexible installation and deployment options. See the guides below for details: -> [!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. -> - If users encounter problems such as missing `lib*.so`, they should uninstall the package they installed and build the binaries manually. +- [Quick Start](https://kvcache-ai.github.io/Mooncake/getting_started/quick-start.html) +- [Build from Source](https://kvcache-ai.github.io/Mooncake/getting_started/build.html) +- [Deployment Guide](https://kvcache-ai.github.io/Mooncake/deployment/mooncake-store-deployment-guide.html) -### Use Docker image -Mooncake supports Docker-based deployment, see [Build Guide](https://kvcache-ai.github.io/Mooncake/getting_started/build.html) in detail. -To produce an image that compiles Mooncake from source, builds the wheel via `scripts/build_wheel.sh`, and installs that wheel inside the container, use `build-wheel.dockerfile`: +### Skills for AI Assistants -```bash -docker build -f docker/mooncake.Dockerfile \ - --build-arg PYTHON_VERSION=3.10 \ - --build-arg EP_TORCH_VERSIONS="2.9.1" \ - -t mooncake:from-source . -``` +Mooncake ships a set of **built-in skills** under [`.claude/skills`](.claude/skills) — reusable, task-focused playbooks that an AI coding assistant (such as Claude Code) invokes automatically when your request matches, or that you can run as a slash command. -The resulting image already has a virtual environment at `/opt/venv` with the freshly built wheel installed. Launch it with GPU/RDMA access as needed, for example: +
+Details -```bash -docker run --gpus all --network host -it mooncake:from-source /bin/bash -``` +| Skill | Description | +|-------|-------------| +| `/mooncake-troubleshoot` | Diagnose Mooncake deployment and runtime issues (services, RDMA, env vars, logs). | +| `/mooncake-ci-local` | Run pre-PR local validation via `scripts/run_ci_test.sh`. | +| `/mooncake-api` | Work with the Mooncake Store, Transfer Engine, and EP/Backend Python APIs. | -> [!NOTE] -> Make sure you build the image from the repository root so that Git metadata and submodules are available inside the build context. - -### Build and use binaries -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)*. -- [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. - -The build and installation steps are as follows: -1. Retrieve source code from GitHub repo - ```bash - git clone https://github.com/kvcache-ai/Mooncake.git - cd Mooncake - ``` - -2. Install dependencies - ```bash - bash dependencies.sh - ``` - -3. Compile Mooncake and examples - ```bash - mkdir build - cd build - cmake .. - make -j - sudo make install # optional, make it ready to be used by vLLM/SGLang - ``` - - -

🛣️ Incoming Milestones

- -- [x] First release of Mooncake and integrate with latest vLLM -- [ ] Share KV caches across multiple serving engines -- [ ] User and developer documentation - -

📦 Open Source Trace

- -```json -{ - "timestamp": 27482, - "input_length": 6955, - "output_length": 52, - "hash_ids": [46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 2353, 2354] -} -{ - "timestamp": 30535, - "input_length": 6472, - "output_length": 26, - "hash_ids": [46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 2366] -} +Install them without cloning the repository via the [Claude Code plugin marketplace](https://code.claude.com/docs/en/plugin-marketplaces): + +```text +/plugin marketplace add kvcache-ai/Mooncake --sparse .claude-plugin +/plugin install mooncake-troubleshoot@mooncake +/plugin install mooncake-ci-local@mooncake +/plugin install mooncake-api@mooncake ``` -The above presents two samples from our trace dataset. The trace includes the timing of request arrivals, the number of input tokens, the number of output tokens, and the remapped block hash. To protect our customers' privacy, we applied several mechanisms to remove user-related information while preserving the dataset's utility for simulated evaluation. More descriptions of the trace (e.g., up to 50% cache hit ratio) can be found in Section 4 of the technical report. -**_Update[Feb 21, 2025]: The updated [traces](FAST25-release/traces) used in our FAST'25 paper have been released! Please refer to the paper's appendix (found [here](FAST25-release/Mooncake-FAST25.pdf)) for more details._** +The `--sparse .claude-plugin` flag fetches only the marketplace catalog, and each plugin is published as a `git-subdir` source, so installing one fetches only that single skill directory — never the whole repo. If you are already working inside a Mooncake checkout, the skills under `.claude/skills/` load automatically with no setup. + +
+ +

📦 Open Source Traces and Tools

+ +We open-source anonymized request traces containing request arrival times, input and output token counts, and remapped block hashes. These traces are designed to support reproducible simulation and evaluation of caching behavior while preserving user privacy. The released traces and related details are available in [FAST25-release](FAST25-release). + +Together with the released traces, we also provide two KV cache analysis tools: a [KV Cache Size Calculator](https://kvcache.ai/tools/kv-cache-size-calculator/) for calculating cache capacity across popular LLM model families, and a [KV Cache Hit Rate Simulator](https://kvcache.ai/tools/kv-cache-hit-rate-simulator/) for analyzing KV cache hit rates and planning cache capacity under different workloads and models. These tools help users better understand KV cache storage costs and caching effectiveness when analyzing or reproducing serving workloads. The tools are open-sourced [here](https://github.com/kvcache-ai/kvcache-blog).

📑 Citation

-Please kindly cite our paper if you find the paper or the traces are useful: +Please kindly cite our papers if you find the papers or the traces are useful: + +```bibtex +@inproceedings{qin2025mooncake, + author = {Ruoyu Qin and Zheming Li and Weiran He and Jialei Cui and Feng Ren and Mingxing Zhang and Yongwei Wu and Weimin Zheng and Xinran Xu}, + title = {Mooncake: Trading More Storage for Less Computation {\textemdash} A {KVCache-centric} Architecture for Serving {LLM} Chatbot}, + booktitle = {23rd USENIX Conference on File and Storage Technologies (FAST 25)}, + year = {2025}, + isbn = {978-1-939133-45-8}, + address = {Santa Clara, CA}, + pages = {155--170}, + url = {https://www.usenix.org/conference/fast25/presentation/qin}, + publisher = {USENIX Association}, + month = {feb}, +} +``` + +
+More ```bibtex +@misc{ren2026tentdeclarativeslicespraying, + title = {TENT: A Declarative Slice Spraying Engine for Performant and Resilient Data Movement in Disaggregated LLM Serving}, + author = {Feng Ren and Ruoyu Qin and Teng Ma and Shangming Cai and Zheng Liu and Chao Lei and Dejiang Zhu and Ke Yang and Zheming Li and Jialei Cui and Weixiao Huang and Yikai Zhao and Yineng Zhang and Hao Wu and Xiang Gao and Yuhao Fu and Jinlei Jiang and Yongwei Wu and Mingxing Zhang}, + year = {2026}, + eprint = {2604.00368}, + archivePrefix = {arXiv}, + primaryClass = {cs.DC}, + url = {https://arxiv.org/abs/2604.00368}, +} + +@article{sun2026survivingpartialrankfailures, + title = {Surviving Partial Rank Failures in Wide Expert-Parallel MoE Inference}, + author = {Xun Sun and Shaoyuan Chen and Pingchuan Ma and Yue Chen and Ziwei Yuan and Zhanhao Cao and Han Han and Shangming Cai and Teng Ma and Xuchun Shang and Xinpeng Zhao and Ke Yang and Junlin Wei and Lianzhi Lin and Yuji Liu and Feng Ren and Haoran Hu and Cheng Wan and Yingdi Shan and Yongwei Wu and Mingxing Zhang}, + year = {2026}, + url = {https://arxiv.org/abs/2605.10670}, +} + @article{qin2025mooncake_tos, author = {Qin Ruoyu and Li Zheming and He Weiran and Cui Jialei and Tang Heyi and Ren Feng and Ma Teng and Cai Shangming and Zhang Yineng and Zhang Mingxing and Wu Yongwei and Zheng Weimin and Xu Xinran}, title = {Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving}, @@ -298,19 +306,6 @@ Please kindly cite our paper if you find the paper or the traces are useful: keywords = {Machine learning system, LLM serving, KVCache}, } -@inproceedings{qin2025mooncake, - author = {Ruoyu Qin and Zheming Li and Weiran He and Jialei Cui and Feng Ren and Mingxing Zhang and Yongwei Wu and Weimin Zheng and Xinran Xu}, - title = {Mooncake: Trading More Storage for Less Computation {\textemdash} A {KVCache-centric} Architecture for Serving {LLM} Chatbot}, - booktitle = {23rd USENIX Conference on File and Storage Technologies (FAST 25)}, - year = {2025}, - isbn = {978-1-939133-45-8}, - address = {Santa Clara, CA}, - pages = {155--170}, - url = {https://www.usenix.org/conference/fast25/presentation/qin}, - publisher = {USENIX Association}, - month = {feb}, -} - @article{qin2024mooncake_arxiv, title = {Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving}, author = {Ruoyu Qin and Zheming Li and Weiran He and Mingxing Zhang and Yongwei Wu and Weimin Zheng and Xinran Xu}, @@ -318,3 +313,5 @@ Please kindly cite our paper if you find the paper or the traces are useful: url = {https://arxiv.org/abs/2407.00079}, } ``` + +
diff --git a/benchmarks/storage_benchmark_v1/__init__.py b/benchmarks/storage_benchmark_v1/__init__.py new file mode 100644 index 0000000000..2973545a0d --- /dev/null +++ b/benchmarks/storage_benchmark_v1/__init__.py @@ -0,0 +1,41 @@ +""" +Mooncake KVCache Storage Benchmark Library +""" + +__version__ = "2.0.0" + +from benchmark import ( + main, + run_benchmark, + print_results, + StorageBenchmark, + TraceReplay, + KVCacheRequest, +) +from storage import Storage, KVKey, KVValue, DiskHashTable, SSDStorage +from layout import ( + KVLayout, MLALayout, KVEntry, MLA_MODEL_CONFIG, get_model_config, create_layout, +) + +__all__ = [ + # Main + 'main', + 'run_benchmark', + 'print_results', + 'StorageBenchmark', + 'TraceReplay', + 'KVCacheRequest', + # Storage + 'Storage', + 'KVKey', + 'KVValue', + 'DiskHashTable', + 'SSDStorage', + # Layout + 'KVLayout', + 'MLALayout', + 'KVEntry', + 'MLA_MODEL_CONFIG', + 'get_model_config', + 'create_layout', +] diff --git a/benchmarks/storage_benchmark_v1/__main__.py b/benchmarks/storage_benchmark_v1/__main__.py new file mode 100644 index 0000000000..2ae8bc8980 --- /dev/null +++ b/benchmarks/storage_benchmark_v1/__main__.py @@ -0,0 +1,11 @@ +""" +Entry point for running the benchmark as a module + +Usage: + python -m storage_benchmark --model=glm5 --scenario=toolagent --max-requests=100 +""" + +from benchmark import main + +if __name__ == '__main__': + main() diff --git a/benchmarks/storage_benchmark_v1/benchmark.py b/benchmarks/storage_benchmark_v1/benchmark.py new file mode 100644 index 0000000000..f840732393 --- /dev/null +++ b/benchmarks/storage_benchmark_v1/benchmark.py @@ -0,0 +1,746 @@ +#!/usr/bin/env python3 +""" +Mooncake KVCache Storage Benchmark + +Complete benchmark tool with CLI interface. +""" + +import argparse +import json +import sys +import time +import statistics +from contextlib import ExitStack +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass +from pathlib import Path +from typing import List, Dict, Any + +from storage import DiskHashTable +from layout import get_model_config, create_layout + + +# ============================================================================ +# Data Structures +# ============================================================================ + +@dataclass +class KVCacheRequest: + """KVCache request from trace""" + timestamp: float + hash_ids: List[int] + input_length: int + output_length: int + + +# ============================================================================ +# Trace Replay +# ============================================================================ + +class TraceReplay: + """Trace replay handler""" + + def __init__(self, trace_path: str): + self.trace_path = trace_path + + def load_all(self) -> List[KVCacheRequest]: + """Load all requests from trace file""" + requests = [] + with open(self.trace_path, 'r', encoding='utf-8') as f: + for line in f: + line = line.strip() + if line: + req = json.loads(line) + requests.append(KVCacheRequest( + timestamp=req.get('timestamp', 0), + hash_ids=req.get('hash_ids', []), + input_length=req.get('input_length', 0), + output_length=req.get('output_length', 0), + )) + return requests + + +# ============================================================================ +# Storage Benchmark +# ============================================================================ + +class StorageBenchmark: + """KVCache storage benchmark + + Processes KVCache requests using layout-generated access patterns. + """ + + def __init__(self, storage_dir: str, model_config: dict, + page_size_tokens: int = 512, + max_pages: int = 100000, + fsync_mode: str = 'none', fsync_batch_size: int = 100): + """Initialize benchmark + + Args: + storage_dir: Directory for storage files + model_config: Model configuration dict + page_size_tokens: Tokens per page (default: 512) + max_pages: Maximum number of pages + fsync_mode: When to fsync ('none', 'batch', 'always', 'end') + fsync_batch_size: Number of writes between fsync in batch mode + """ + self.model_config = model_config + self.layout = create_layout(model_config, page_size_tokens) + self.page_size_bytes = self.layout.value_size_bytes + + # Initialize storage + self.storage = DiskHashTable( + storage_dir=storage_dir, + page_size=self.page_size_bytes, + max_pages=max_pages, + fsync_mode=fsync_mode, + fsync_batch_size=fsync_batch_size + ) + + # Statistics + self.stats = { + 'total_requests': 0, + 'total_tokens': 0, + 'read_pages': 0, + 'write_pages': 0, + 'page_hits': 0, + 'request_io_latencies_ms': [], + 'request_wall_latencies_ms': [], + } + + def process_request(self, req: KVCacheRequest) -> float: + """Process a KVCache request + + Args: + req: KVCache request + + Returns: + Total latency in milliseconds + """ + self.stats['total_requests'] += 1 + self.stats['total_tokens'] += req.input_length + req.output_length + + request_start = time.perf_counter() + io_latency_ms = 0.0 + + # Process each access requirement from layout + for access in self.layout.get_operations(req): + if self.storage.exists(access.page_id): + # Page exists, perform READ + io_latency_ms += self.storage.read( + access.page_id, + offset_in_page=access.offset_in_page, + length=access.length + ) + self.stats['read_pages'] += 1 + self.stats['page_hits'] += 1 + else: + # Page doesn't exist, perform WRITE + io_latency_ms += self.storage.write( + access.page_id, + offset_in_page=access.offset_in_page, + length=access.length + ) + self.stats['write_pages'] += 1 + + wall_latency_ms = (time.perf_counter() - request_start) * 1000.0 + self.stats['request_io_latencies_ms'].append(io_latency_ms) + self.stats['request_wall_latencies_ms'].append(wall_latency_ms) + return io_latency_ms + + def get_stats(self) -> Dict: + """Get statistics""" + storage_stats = self.storage.get_stats() + request_io_latencies = self.stats['request_io_latencies_ms'] + request_wall_latencies = self.stats['request_wall_latencies_ms'] + + total_pages = self.stats['read_pages'] + self.stats['write_pages'] + + return { + 'total_requests': self.stats['total_requests'], + 'total_tokens': self.stats['total_tokens'], + 'total_pages': total_pages, + 'read_pages': self.stats['read_pages'], + 'write_pages': self.stats['write_pages'], + 'page_hits': self.stats['page_hits'], + 'page_hit_rate': self.stats['read_pages'] / total_pages if total_pages > 0 else 0, + 'write_ratio': self.stats['write_pages'] / total_pages if total_pages > 0 else 0, + 'request_io_latency': latency_stats(request_io_latencies), + 'request_wall_latency': latency_stats(request_wall_latencies), + 'storage': storage_stats, + } + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() + return False + + def close(self, force_sync: bool = True): + self.storage.close(force_sync=force_sync) + + +# ============================================================================ +# Benchmark Runner +# ============================================================================ + +def get_max_page_id(requests: List[KVCacheRequest]) -> int: + max_id = 0 + for req in requests: + if req.hash_ids: + max_id = max(max_id, max(req.hash_ids)) + return max_id + + +def parse_csv_floats(value: str) -> List[float]: + return [float(item.strip()) for item in value.split(',') if item.strip()] + + +def wait_for_replay_time(req: KVCacheRequest, base_timestamp: float, + start_time: float, replay_scale: float): + if replay_scale <= 0 or req.timestamp == 0: + return + target_time = (start_time + + max(0.0, req.timestamp - base_timestamp) / + (1000.0 * replay_scale)) + delay = target_time - time.perf_counter() + if delay > 0: + time.sleep(delay) + + +def latency_stats(values: List[float]) -> Dict[str, float]: + if not values: + return {'avg_ms': 0, 'p50_ms': 0, 'p95_ms': 0, 'p99_ms': 0} + + sorted_values = sorted(values) + + def get_percentile(p: float) -> float: + if len(sorted_values) == 1: + return sorted_values[0] + rank = (len(sorted_values) - 1) * p + lower = int(rank) + upper = min(lower + 1, len(sorted_values) - 1) + weight = rank - lower + return (sorted_values[lower] * (1.0 - weight) + + sorted_values[upper] * weight) + + return { + 'avg_ms': statistics.mean(values), + 'p50_ms': get_percentile(0.50), + 'p95_ms': get_percentile(0.95), + 'p99_ms': get_percentile(0.99), + } + + +def snapshot_thread_stats(benchmark: StorageBenchmark) -> Dict[str, Any]: + storage = benchmark.storage + total_pages = benchmark.stats['read_pages'] + benchmark.stats['write_pages'] + return { + 'total_requests': benchmark.stats['total_requests'], + 'total_tokens': benchmark.stats['total_tokens'], + 'read_pages': benchmark.stats['read_pages'], + 'write_pages': benchmark.stats['write_pages'], + 'page_hits': benchmark.stats['page_hits'], + 'request_io_latencies_ms': list( + benchmark.stats['request_io_latencies_ms'] + ), + 'request_wall_latencies_ms': list( + benchmark.stats['request_wall_latencies_ms'] + ), + 'read_bytes': storage.stats['read_bytes'], + 'write_bytes': storage.stats['write_bytes'], + 'read_time_s': storage.stats['read_time_s'], + 'write_time_s': storage.stats['write_time_s'], + 'read_latencies_ms': list(storage.stats['read_latencies_ms']), + 'write_latencies_ms': list(storage.stats['write_latencies_ms']), + 'sync_count': storage.stats['sync_count'], + 'max_pages': storage.max_pages, + 'written_pages': len(storage._written_pages), + 'total_pages': total_pages, + } + + +def aggregate_thread_stats(thread_stats: List[Dict[str, Any]]) -> Dict: + total_requests = sum(s['total_requests'] for s in thread_stats) + total_tokens = sum(s['total_tokens'] for s in thread_stats) + read_pages = sum(s['read_pages'] for s in thread_stats) + write_pages = sum(s['write_pages'] for s in thread_stats) + page_hits = sum(s['page_hits'] for s in thread_stats) + total_pages = read_pages + write_pages + + request_io_latencies = [] + request_wall_latencies = [] + read_latencies = [] + write_latencies = [] + for stats in thread_stats: + request_io_latencies.extend(stats['request_io_latencies_ms']) + request_wall_latencies.extend(stats['request_wall_latencies_ms']) + read_latencies.extend(stats['read_latencies_ms']) + write_latencies.extend(stats['write_latencies_ms']) + + read_bytes = sum(s['read_bytes'] for s in thread_stats) + write_bytes = sum(s['write_bytes'] for s in thread_stats) + read_time = sum(s['read_time_s'] for s in thread_stats) + write_time = sum(s['write_time_s'] for s in thread_stats) + + return { + 'total_requests': total_requests, + 'total_tokens': total_tokens, + 'total_pages': total_pages, + 'read_pages': read_pages, + 'write_pages': write_pages, + 'page_hits': page_hits, + 'page_hit_rate': read_pages / total_pages if total_pages > 0 else 0, + 'write_ratio': write_pages / total_pages if total_pages > 0 else 0, + 'request_io_latency': latency_stats(request_io_latencies), + 'request_wall_latency': latency_stats(request_wall_latencies), + 'storage': { + 'read': { + 'count': read_pages, + 'mb': read_bytes / 1024 / 1024, + 'time_s': read_time, + **latency_stats(read_latencies), + }, + 'write': { + 'count': write_pages, + 'mb': write_bytes / 1024 / 1024, + 'time_s': write_time, + **latency_stats(write_latencies), + }, + 'sync_count': sum(s['sync_count'] for s in thread_stats), + 'max_pages': sum(s['max_pages'] for s in thread_stats), + 'written_pages': sum(s['written_pages'] for s in thread_stats), + 'page_hits': page_hits, + 'page_misses': write_pages, + }, + } + + +def print_progress(done: int, total: int, start_time: float, + stats: Dict, req: KVCacheRequest = None, + suffix: str = ""): + elapsed = time.perf_counter() - start_time + qps = done / elapsed if elapsed > 0 else 0 + storage = stats.get('storage', {}) + read_stats = storage.get('read', {}) + write_stats = storage.get('write', {}) + read_time = read_stats.get('time_s', 0) + write_time = write_stats.get('time_s', 0) + read_mbps = read_stats.get('mb', 0) / read_time if read_time > 0 else 0 + write_mbps = write_stats.get('mb', 0) / write_time if write_time > 0 else 0 + + if req is None: + req_info = "" + else: + req_info = (f" ids={len(req.hash_ids):3d} " + f"tokens={req.input_length + req.output_length:6d} |") + + print(f" [{done:5d}/{total}]{req_info} QPS={qps:7.2f} | " + f"R={stats['read_pages']:6d} " + f"({read_stats.get('avg_ms', 0):6.2f}ms, {read_mbps:6.1f}MB/s) | " + f"W={stats['write_pages']:6d} " + f"({write_stats.get('avg_ms', 0):6.2f}ms, {write_mbps:6.1f}MB/s)" + f"{suffix}") + + +def should_print_progress(done: int, total: int, progress_interval: int) -> bool: + if done >= total: + return True + return progress_interval > 0 and done % progress_interval == 0 + + +def run_single_thread(benchmark: StorageBenchmark, + requests: List[KVCacheRequest], + replay_scale: float, + progress_interval: int) -> Dict[str, Any]: + start_time = time.perf_counter() + base_timestamp = requests[0].timestamp if requests else 0 + completed = 0 + + for req in requests: + wait_for_replay_time(req, base_timestamp, start_time, replay_scale) + benchmark.process_request(req) + completed += 1 + if should_print_progress(completed, len(requests), progress_interval): + print_progress(completed, len(requests), start_time, + benchmark.get_stats(), req) + + return { + 'completed': completed, + 'elapsed': time.perf_counter() - start_time, + 'stats': benchmark.get_stats(), + } + + +def run_multi_thread(benchmarks: List[StorageBenchmark], + requests: List[KVCacheRequest], + replay_scale: float) -> Dict[str, Any]: + start_time = time.perf_counter() + base_timestamp = requests[0].timestamp if requests else 0 + total_requests = len(requests) * len(benchmarks) + completed = 0 + + def run_worker(thread_id: int): + benchmark = benchmarks[thread_id] + for req in requests: + wait_for_replay_time(req, base_timestamp, start_time, replay_scale) + benchmark.process_request(req) + return snapshot_thread_stats(benchmark) + + thread_stats = [] + with ThreadPoolExecutor(max_workers=len(benchmarks)) as executor: + futures = [ + executor.submit(run_worker, thread_id) + for thread_id in range(len(benchmarks)) + ] + for future in as_completed(futures): + worker_stats = future.result() + thread_stats.append(worker_stats) + completed += worker_stats['total_requests'] + print_progress(completed, total_requests, start_time, + aggregate_thread_stats(thread_stats), + suffix=" | completed worker") + + return { + 'completed': completed, + 'elapsed': time.perf_counter() - start_time, + 'stats': aggregate_thread_stats(thread_stats), + } + + +def run_benchmark(trace_path: str, storage_dir: str, model_config: dict, + max_requests: int = None, max_pages: int = None, + page_size_tokens: int = 512, + fsync_mode: str = 'none', fsync_batch_size: int = 100, + threads: int = 1, replay_scale: float = 0.0, + progress_interval: int = 100) -> Dict: + """Run benchmark + + Args: + trace_path: Trace file path + storage_dir: Storage directory + model_config: Model configuration + max_requests: Maximum number of requests (None = all) + max_pages: Maximum number of pages (None = auto-calculate) + page_size_tokens: Tokens per page + fsync_mode: When to fsync + fsync_batch_size: Number of writes between fsync + threads: Benchmark client worker threads + replay_scale: Timestamp replay multiplier; 0 runs unpaced + progress_interval: Print progress every N requests; 0 disables progress + + Returns: + Benchmark results dictionary + """ + print(f"\n{'='*80}") + print(f"Running: {Path(trace_path).name}") + print(f"Model: {model_config['name']}") + print(f"Layers: {model_config['num_layers']}") + print(f"Page size: {page_size_tokens} tokens") + print(f"Threads: {threads}") + print(f"Fast-forward: {replay_scale:g}x" if replay_scale > 0 else "Fast-forward: unpaced") + print(f"{'='*80}") + + # Load trace + replay = TraceReplay(trace_path) + requests = replay.load_all() + + if max_requests: + requests = requests[:max_requests] + + print(f"Loaded {len(requests)} requests") + + # Find max page_id from trace + max_page_id = get_max_page_id(requests) + max_pages_needed = max_page_id + 1 # page_id is 0-based + + # Create layout to get page size + layout = create_layout(model_config, page_size_tokens) + page_size_bytes = layout.value_size_bytes + + # Determine max_pages + if max_pages is None: + # Use max page_id from trace + max_pages = max_pages_needed + # Round up to next thousand for cleaner numbers + max_pages = ((max_pages + 999) // 1000) * 1000 + else: + # User specified max_pages + pass + + max_size_gb = max_pages * page_size_bytes / (1024**3) + trace_size_gb = max_pages_needed * page_size_bytes / (1024**3) + + print(f"\n[Storage Configuration]") + print(f" Max page_id in trace: {max_page_id:,}") + print(f" Pages needed (trace): {max_pages_needed:,}") + print(f" Trace storage size: {trace_size_gb:.2f} GB") + print(f" Max pages configured: {max_pages:,}") + if threads > 1: + print(f" Max pages across threads: {max_pages * threads:,}") + print(f" Max storage available: {max_size_gb:.2f} GB") + if threads > 1: + print(f" Max storage across threads: {max_size_gb * threads:.2f} GB") + + if max_pages_needed > max_pages: + shortfall = max_pages_needed - max_pages + shortfall_gb = shortfall * page_size_bytes / (1024**3) + compression_ratio = max_pages / max_pages_needed + print(f"\n ⚠️ Storage insufficient: {shortfall:,} pages shortfall ({shortfall_gb:.2f} GB)") + print(f" ⚠️ Consider increasing --max-pages to at least {max_pages_needed:,} for full simulation") + else: + surplus = max_pages - max_pages_needed + surplus_pct = (surplus / max_pages) * 100 if max_pages > 0 else 0 + print(f" ✓ Direct mapping: all {max_pages_needed:,} logical pages uniquely mapped") + + try: + if threads <= 1: + with StorageBenchmark( + storage_dir=storage_dir, + model_config=model_config, + page_size_tokens=page_size_tokens, + max_pages=max_pages, + fsync_mode=fsync_mode, + fsync_batch_size=fsync_batch_size + ) as benchmark: + result = run_single_thread(benchmark, requests, replay_scale, + progress_interval) + else: + with ExitStack() as stack: + benchmarks = [ + stack.enter_context(StorageBenchmark( + storage_dir=str(Path(storage_dir) / f"thread_{thread_id}"), + model_config=model_config, + page_size_tokens=page_size_tokens, + max_pages=max_pages, + fsync_mode=fsync_mode, + fsync_batch_size=fsync_batch_size + )) + for thread_id in range(threads) + ] + result = run_multi_thread(benchmarks, requests, replay_scale) + except KeyboardInterrupt: + print(f"\n\n{'='*80}") + print(f"Interrupted! Showing partial results:") + print(f"{'='*80}") + result = result if 'result' in locals() else { + 'completed': 0, + 'elapsed': 0, + 'stats': {}, + } + print_results([{ + 'trace_file': Path(trace_path).name, + 'total_requests': result['completed'], + 'io_time_s': result['elapsed'], + 'requests_per_second': ( + result['completed'] / result['elapsed'] + if result['elapsed'] > 0 else 0 + ), + 'model': model_config['name'], + 'fsync_mode': fsync_mode, + 'threads': threads, + 'replay_scale': replay_scale, + **result['stats'], + }]) + sys.exit(0) + + return { + 'trace_file': Path(trace_path).name, + 'total_requests': result['completed'], + 'io_time_s': result['elapsed'], + 'requests_per_second': ( + result['completed'] / result['elapsed'] if result['elapsed'] > 0 else 0 + ), + 'model': model_config['name'], + 'fsync_mode': fsync_mode, + 'threads': threads, + 'replay_scale': replay_scale, + **result['stats'], + } + + +# ============================================================================ +# Output Formatting +# ============================================================================ + +def format_storage_stats(stats: Dict, title: str = "Storage"): + """Format storage statistics with clear read/write separation""" + storage = stats.get('storage', {}) + read_stats = storage.get('read', {}) + write_stats = storage.get('write', {}) + request_wall = stats.get('request_wall_latency', {}) + request_io = stats.get('request_io_latency', {}) + + output = [] + output.append(f"\n[{title}]") + + # General info + output.append(f"\n[General]") + output.append(f" Model: {stats.get('model', 'N/A')}") + output.append(f" Threads: {stats.get('threads', 1)}") + replay_scale = stats.get('replay_scale', 0) + output.append(f" Fast-forward: {f'{replay_scale:g}x' if replay_scale else 'unpaced'}") + output.append(f" Requests: {stats.get('total_requests', 0):,}") + output.append(f" Tokens: {stats.get('total_tokens', 0):,}") + output.append(f" Total I/O Time: {stats.get('io_time_s', 0):.3f} s") + output.append(f" QPS: {stats.get('requests_per_second', 0):.2f}") + output.append(f" Hit Rate: {stats.get('page_hit_rate', 0):.2%}") + + # Request Stats + output.append(f"\n[Request Wall Latency]") + output.append(f" Avg: {request_wall.get('avg_ms', 0):.3f} ms") + output.append(f" P50: {request_wall.get('p50_ms', 0):.3f} ms") + output.append(f" P95: {request_wall.get('p95_ms', 0):.3f} ms") + output.append(f" P99: {request_wall.get('p99_ms', 0):.3f} ms") + + output.append(f"\n[Request Storage I/O Latency]") + output.append(f" Avg: {request_io.get('avg_ms', 0):.3f} ms") + output.append(f" P50: {request_io.get('p50_ms', 0):.3f} ms") + output.append(f" P95: {request_io.get('p95_ms', 0):.3f} ms") + output.append(f" P99: {request_io.get('p99_ms', 0):.3f} ms") + + # Read Stats + output.append(f"\n[Read Operations]") + output.append(f" Count: {read_stats.get('count', 0):,}") + output.append(f" Data Volume: {read_stats.get('mb', 0):.2f} MB") + read_time = read_stats.get('time_s', 0) + read_mbps = read_stats.get('mb', 0) / read_time if read_time > 0 else 0 + output.append(f" Total Time: {read_time:.3f} s") + output.append(f" Bandwidth: {read_mbps:.2f} MB/s") + output.append(f" Latency:") + output.append(f" Avg: {read_stats.get('avg_ms', 0):.3f} ms") + output.append(f" P50: {read_stats.get('p50_ms', 0):.3f} ms") + output.append(f" P95: {read_stats.get('p95_ms', 0):.3f} ms") + output.append(f" P99: {read_stats.get('p99_ms', 0):.3f} ms") + + # Write Stats + output.append(f"\n[Write Operations]") + output.append(f" Count: {write_stats.get('count', 0):,}") + output.append(f" Data Volume: {write_stats.get('mb', 0):.2f} MB") + write_time = write_stats.get('time_s', 0) + write_mbps = write_stats.get('mb', 0) / write_time if write_time > 0 else 0 + output.append(f" Total Time: {write_time:.3f} s") + output.append(f" Bandwidth: {write_mbps:.2f} MB/s") + output.append(f" Latency:") + output.append(f" Avg: {write_stats.get('avg_ms', 0):.3f} ms") + output.append(f" P50: {write_stats.get('p50_ms', 0):.3f} ms") + output.append(f" P95: {write_stats.get('p95_ms', 0):.3f} ms") + output.append(f" P99: {write_stats.get('p99_ms', 0):.3f} ms") + + # Storage Info + output.append(f"\n[Storage Info]") + output.append(f" Max Pages: {storage.get('max_pages', 0):,}") + output.append(f" Written Pages: {storage.get('written_pages', 0):,}") + output.append(f" Sync Count: {storage.get('sync_count', 0):,}") + + return "\n".join(output) + + +def print_results(results: List[Dict]): + """Print 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(format_storage_stats(r)) + + +# ============================================================================ +# CLI Entry Point +# ============================================================================ + +def main(): + """Main entry point""" + parser = argparse.ArgumentParser( + description='Mooncake KVCache Storage Benchmark', + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + + 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, default='glm5', choices=['glm5', 'kimi-k2.6'], + help='Model preset') + parser.add_argument('--page-size-tokens', type=int, default=512, + help='Page size in tokens (default: 512)') + parser.add_argument('--max-requests', type=int, default=None, + help='Maximum number of requests') + parser.add_argument('--max-pages', type=int, default=2000, + help='Maximum number of pages') + parser.add_argument('--fsync-mode', type=str, choices=['batch', 'always', 'end', 'none'], + default='none', help='When to fsync') + parser.add_argument('--fsync-batch-size', type=int, default=100, + help='Number of writes between fsync') + parser.add_argument('--threads', type=int, default=1, + help='Number of benchmark client worker threads') + parser.add_argument('--replay-scales', type=str, default='0', + help='Comma-separated trace fast-forward speeds; 0 means unpaced') + parser.add_argument('--progress-interval', type=int, default=100, + help='Print progress every N requests; 0 disables per-request progress') + + args = parser.parse_args() + if args.threads < 1: + parser.error('--threads must be at least 1') + if args.progress_interval < 0: + parser.error('--progress-interval must be non-negative') + + print(f"\n{'='*80}") + print(f"{'Mooncake KVCache Storage Benchmark':^80}") + print(f"{'='*80}") + + model_config = get_model_config(args.model) + print(f"Model: {args.model} ({model_config['num_layers']} layers)") + replay_scales = parse_csv_floats(args.replay_scales) + if not replay_scales: + parser.error('--replay-scales must include at least one value') + if any(scale < 0 for scale in replay_scales): + parser.error('--replay-scales values must be non-negative') + + # Determine 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 = [] + use_scale_subdirs = len(replay_scales) > 1 or replay_scales[0] != 0 + for scenario in scenarios: + trace_path = Path(args.trace_dir) / trace_files[scenario] + if trace_path.exists(): + for replay_scale in replay_scales: + run_dir = Path(args.storage_dir) / scenario + if use_scale_subdirs: + run_dir = run_dir / f"replay_{replay_scale:g}x" + result = run_benchmark( + str(trace_path), + str(run_dir), + model_config, + args.max_requests, + args.max_pages, + args.page_size_tokens, + args.fsync_mode, + args.fsync_batch_size, + args.threads, + replay_scale, + args.progress_interval + ) + results.append(result) + else: + print(f"Warning: Trace file not found: {trace_path}") + + # Print results + if results: + print_results(results) + else: + print("Error: No trace files were successfully processed.", file=sys.stderr) + sys.exit(1) + +if __name__ == '__main__': + main() diff --git a/benchmarks/storage_benchmark_v1/doc/README.md b/benchmarks/storage_benchmark_v1/doc/README.md new file mode 100644 index 0000000000..8b53be660a --- /dev/null +++ b/benchmarks/storage_benchmark_v1/doc/README.md @@ -0,0 +1,192 @@ +# KVCache Storage Benchmark v1 + +## Overview + +The KVCache Storage Benchmark is a tool for evaluating storage performance of KVCache workloads. It simulates real-world cache access patterns using trace replay and measures storage I/O performance with detailed statistics. + +## Usage + +### Basic Usage + +```bash +cd benchmarks/storage_benchmark_v1 +python benchmark.py --scenario conversation \ + --trace-dir /path/to/Mooncake/FAST25-release/traces \ + --storage-dir /path/to/test/drive +``` + +### Command Line Options + +| Option | Default | Description | +|--------|---------|-------------| +| `--trace-dir` | `../../FAST25-release/traces` | Directory containing trace files | +| `--scenario` | `toolagent` | Test scenario: `conversation`, `synthetic`, `toolagent`, or `all` | +| `--storage-dir` | `/tmp/mooncake_bench` | Directory for storage files | +| `--model` | `glm5` | Model preset: `glm5` or `kimi-k2.6` | +| `--page-size-tokens` | `512` | Page size in tokens | +| `--max-requests` | `None` | Maximum number of requests to process | +| `--max-pages` | `2000` | Maximum number of pages (creates modulo mapping if trace is larger) | +| `--fsync-mode` | `none` | When to fsync: `none`, `batch`, `always`, or `end` | +| `--fsync-batch-size` | `100` | Number of writes between fsync in batch mode | +| `--threads` | `1` | Number of benchmark client worker threads | +| `--replay-scales` | `0` | Comma-separated trace fast-forward speeds; `0` means unpaced | +| `--progress-interval` | `100` | Print progress every N requests; `0` disables per-request progress | + +### Replay Scale + +Use `--replay-scales` to run the same trace at different fast-forward speeds: + +```bash +python benchmark.py --scenario toolagent \ + --trace-dir /path/to/Mooncake/FAST25-release/traces \ + --storage-dir /path/to/test/drive \ + --replay-scales 1,2,4,8 +``` + +For example, `2` means 2x fast-forward and `8` means 8x fast-forward. `0` +preserves the old unpaced behavior. + +### Client Threads + +Use `--threads` to add benchmark client worker threads: + +```bash +python benchmark.py --scenario toolagent \ + --trace-dir /path/to/Mooncake/FAST25-release/traces \ + --storage-dir /path/to/test/drive \ + --threads 4 +``` + +With `--threads > 1`, each benchmark client thread uses an independent storage +file under `thread_N/data.bin`, similar to running multiple clients at the same +time. Final results aggregate the per-thread counters and latency samples. For +strict single-client trace-order read/write and hit-rate accounting, use +`--threads 1`. + +## Output Format + +### Progress Output + +During execution, progress is printed every `--progress-interval` requests and +at the end of the run: + +``` +[ 10/12031] ids= 35 tokens= 18060 | QPS= 2.45 | R= 36 ( 22.01ms, 2435.2MB/s) | W= 963 ( 19.35ms, 2770.1MB/s) +``` + +Fields: +- `[N/Total]`: Current request progress +- `ids=N`: Number of hash_ids in this request +- `tokens=N`: Total tokens (input + output) +- `QPS=X`: Queries per second (overall) +- `R=N (latency, bandwidth)`: Read count, average latency, bandwidth +- `W=N (latency, bandwidth)`: Write count, average latency, bandwidth + +### Final Results + +``` +================================================================================ + [1/1] toolagent_trace.jsonl +================================================================================ + +[General] + Model: glm5 + Threads: 1 + Fast-forward: unpaced + Requests: 12031 + Tokens: 123456789 + Total I/O Time: 245.123 s + QPS: 49.07 + Hit Rate: 3.25% + +[Request Wall Latency] + Avg: 20.912 ms + P50: 19.654 ms + P95: 28.123 ms + P99: 34.987 ms + +[Request Storage I/O Latency] + Avg: 20.312 ms + P50: 18.987 ms + P95: 27.456 ms + P99: 33.210 ms + +[Read Operations] + Count: 390 + Data Volume: 20919.62 MB + Total Time: 8.590 s + Bandwidth: 2435.67 MB/s + Latency: + Avg: 22.032 ms + P50: 21.456 ms + P95: 28.912 ms + P99: 35.234 ms + +[Write Operations] + Count: 11641 + Data Volume: 654321.45 MB + Total Time: 236.533 s + Bandwidth: 2765.89 MB/s + Latency: + Avg: 20.312 ms + P50: 19.876 ms + P95: 25.123 ms + P99: 31.456 ms + +[Storage Info] + Max Pages: 2000 + Written Pages: 2000 + Sync Count: 0 +``` + +`Request Wall Latency` measures the benchmark client's wall-clock time spent +processing a request after replay pacing. `Request Storage I/O Latency` is the +sum of the request's page read/write latencies. Read/write operation latency is +reported per page operation. Percentile values use linear interpolation. + +## Measurement Notes + +- The default `--fsync-mode none` measures page-cache-backed write behavior. It + does not represent durable write latency. Use `--fsync-mode always`, `batch`, + or `end` when persistence cost is part of the benchmark target. +- `pread`/`pwrite` latency is measured from user space, so it can include page + cache effects, OS scheduling, and Python benchmark-client overhead. Treat the + reported latency as an observed storage-path latency, not raw device service + time. +- With `--threads > 1`, each thread replays the full trace as an independent + benchmark client with its own storage file. This is a multi-client drive test, + not parallel execution of one trace stream. +- For publication-quality numbers, use a fixed machine and storage device, + clear or isolate benchmark storage directories between runs, disable + per-request progress output with `--progress-interval 0`, and run multiple + trials before reporting stable statistics. + +## Modulo Mapping + +When the trace requires more pages than `--max-pages`, modulo mapping is enabled: + +``` +physical_page_id = logical_page_id % max_pages +``` + +This allows simulating large traces (millions of pages) with limited storage (thousands of pages). The first-hit/read-then-write logic is preserved by tracking written logical pages in memory. + +**Example**: With `--max-pages 2000`, logical page IDs 0-1999 map directly to physical pages 0-1999. Logical page 2000 maps to physical page 0, logical page 2001 maps to physical page 1, etc. + +## Graceful Interruption + +Press `Ctrl-C` at any time to stop the benchmark and view partial results. The output will display all statistics collected up to the interruption point, using the same format as final results. + +## Troubleshooting + +### Insufficient Storage Warning +``` +⚠️ Modulo mapping ENABLED (limited storage) +⚠️ Storage insufficient: 14,257,620 pages shortfall (762.34 GB) +``` +Increase `--max-pages` to reduce modulo mapping effects. + +### Low Bandwidth or High Latency +- Check fsync mode (`--fsync-mode none` for best performance) +- Verify disk performance with `fio` or `dd` +- Verify storage device health diff --git a/benchmarks/storage_benchmark_v1/layout/__init__.py b/benchmarks/storage_benchmark_v1/layout/__init__.py new file mode 100644 index 0000000000..5d00a78909 --- /dev/null +++ b/benchmarks/storage_benchmark_v1/layout/__init__.py @@ -0,0 +1,17 @@ +""" +KVCache Layout Module + +Provides layout interface and implementations for different model architectures. +""" + +from .interface import KVLayout, StorageAccess +from .mla import MLALayout, MLA_MODEL_CONFIG, get_model_config, create_layout + +__all__ = [ + 'KVLayout', + 'StorageAccess', + 'MLALayout', + 'MLA_MODEL_CONFIG', + 'get_model_config', + 'create_layout', +] diff --git a/benchmarks/storage_benchmark_v1/layout/interface.py b/benchmarks/storage_benchmark_v1/layout/interface.py new file mode 100644 index 0000000000..569bc10426 --- /dev/null +++ b/benchmarks/storage_benchmark_v1/layout/interface.py @@ -0,0 +1,55 @@ +""" +KVCache Layout Interface + +Defines the abstract interface for KVCache storage layouts. +Layout layer converts normalized requests into storage access requirements. +""" + +from abc import ABC, abstractmethod +from typing import Iterator, Any +from dataclasses import dataclass + + +@dataclass +class StorageAccess: + """Storage access requirement + + Represents a need to access a page. Whether to READ or WRITE is determined + by the upper layer based on whether the page already exists. + """ + page_id: int # Page ID (hash_id from trace) + offset_in_page: int = 0 # Offset within the page (default: 0) + length: int = None # Number of bytes (default: entire page) + + def __repr__(self): + if self.offset_in_page == 0 and self.length is None: + return f"Access(page_id={self.page_id})" + else: + return f"Access(page_id={self.page_id}, offset={self.offset_in_page}, length={self.length})" + + +class KVLayout(ABC): + """Abstract interface for KVCache storage layout + + Converts normalized KVCache requests into storage access requirements. + + Request format: + - hash_ids: List[int] - chunk/page identifiers + - input_length: int - input token count + - output_length: int - output token count + + Output: + - Iterator of StorageAccess (page access requirements) + """ + + @abstractmethod + def get_operations(self, request: Any) -> Iterator[StorageAccess]: + """Generate storage access requirements for a request + + Args: + request: KVCache request with hash_ids, input_length, output_length + + Yields: + StorageAccess: Page access requirements (READ vs WRITE decided by upper layer) + """ + pass diff --git a/benchmarks/storage_benchmark_v1/layout/mla.py b/benchmarks/storage_benchmark_v1/layout/mla.py new file mode 100644 index 0000000000..6d7a38646d --- /dev/null +++ b/benchmarks/storage_benchmark_v1/layout/mla.py @@ -0,0 +1,174 @@ +""" +MLA (Multi-head Latent Attention) KVCache Layout + +Implements the KVLayout interface for MLA architecture. +""" + +from typing import Iterator, Any + +from .interface import KVLayout, StorageAccess + + +# ============================================================================ +# MLA Model Configurations +# ============================================================================ + +# MLA Model configurations +# Source: https://kvcache.ai/tools/kv-cache-calculator/ +# MLA architecture: hash_id -> {layer_0: [pages], layer_1: [pages], ...} +MLA_MODEL_CONFIG = { + # GLM-5: 78 layers, 64 tokens/page + # KV: 78 layers × 64 tokens × (512+64+128) × 2 = 90,112 bytes/page + # Per token: 1,408 bytes + "glm5": { + "name": "glm5", + "num_layers": 78, + "kv_lora_rank": 512, + "qk_rope_head_dim": 64, + "index_head_dim": 128, + "kv_precision_bytes": 2, # BF16 + "indexer_precision_bytes": 2, # BF16 + }, + + # Kimi-K2.6: 61 layers, 64 tokens/page + # KV: 61 layers × 64 tokens × (512+64) × 2 = 73,728 bytes/page + # Per token: 1,152 bytes + "kimi-k2.6": { + "name": "kimi-k2.6", + "num_layers": 61, + "kv_lora_rank": 512, + "qk_rope_head_dim": 64, + "index_head_dim": 0, # Kimi doesn't use separate indexer + "kv_precision_bytes": 2, # BF16 + "indexer_precision_bytes": 0, + }, +} + + +class MLALayout(KVLayout): + """MLA (Multi-head Latent Attention) KVCache layout + + MLA Architecture: + - Each hash_id corresponds to a 512-token chunk + - Each hash_id maps to ONE complete entry containing all layers + - Entry contains KV + Indexer data for all layers for 512 tokens + - Value size is fixed per hash_id (includes all layers) + + Value Size Calculation (per hash_id entry): + per_layer_size = 512 × (kv_lora_rank + qk_rope_head_dim + index_head_dim) × precision_bytes + value_size = per_layer_size × num_layers + + For GLM-5 with 512 tokens per entry: + per_layer_size = 512 × (512 + 64 + 128) × 2 = 720,896 bytes + value_size = 720,896 × 78 = 56,229,888 bytes = 53.6 MiB + + Key Pattern: hash_id → single entry (all layers included) + Total Keys = len(hash_ids) + + Used in: GLM-5, Kimi-K2.6 + """ + + def __init__(self, num_layers: int, kv_lora_rank: int, qk_rope_head_dim: int, + index_head_dim: int, precision_bytes: int, page_size_tokens: int = 512): + """Initialize MLA layout + + Args: + num_layers: Number of transformer layers + kv_lora_rank: KV LoRA rank dimension + qk_rope_head_dim: QK rope head dimension + index_head_dim: Indexer head dimension + precision_bytes: Precision in bytes (BF16=2, INT8=1, INT4=0.5) + page_size_tokens: Tokens per page (default: 512) + """ + self.num_layers = num_layers + self.kv_lora_rank = kv_lora_rank + self.qk_rope_head_dim = qk_rope_head_dim + self.index_head_dim = index_head_dim + self.precision_bytes = precision_bytes + self.page_size_tokens = page_size_tokens + + # Calculate fixed value size per entry (per hash_id) + # Each entry contains KV + Indexer for all layers for page_size_tokens + # Per layer: page_size_tokens × (kv_lora_rank + qk_rope_head_dim + index_head_dim) × precision_bytes + # Total: per_layer_size × num_layers + per_layer_size = page_size_tokens * (kv_lora_rank + qk_rope_head_dim + index_head_dim) * precision_bytes + self.value_size_bytes = per_layer_size * num_layers + + # Store page_size for backward compatibility + self.page_size = self.value_size_bytes + + def get_operations(self, request: Any) -> Iterator[StorageAccess]: + """Generate storage access requirements for a request + + For MLA architecture: + - Each hash_id corresponds to one complete page (512 tokens, all layers) + - Generate one access requirement per hash_id + + Args: + request: KVCache request with hash_ids, input_length, output_length + + Yields: + StorageAccess: Page access requirements + """ + for hash_id in request.hash_ids: + yield StorageAccess( + page_id=hash_id, + offset_in_page=0, + length=self.value_size_bytes + ) + + +# ============================================================================ +# Utility Functions +# ============================================================================ + +def get_model_config(model_name: str) -> dict: + """Get MLA model configuration by name + + Args: + model_name: Model identifier (e.g., 'glm5', 'kimi-k2.6') + + Returns: + dict: Model configuration + + Raises: + KeyError: If model name is not found + """ + if model_name not in MLA_MODEL_CONFIG: + available = ", ".join(MLA_MODEL_CONFIG.keys()) + raise KeyError(f"Unknown model: {model_name}. Available: {available}") + return MLA_MODEL_CONFIG[model_name].copy() + + +def create_layout(model_config: dict, page_size_tokens: int = 64) -> MLALayout: + """Create an MLALayout instance from model configuration + + Args: + model_config: Model configuration dictionary with fields: + - num_layers: Number of transformer layers + - kv_lora_rank: KV LoRA rank dimension + - qk_rope_head_dim: QK rope head dimension + - index_head_dim: Indexer head dimension + - kv_precision_bytes: Precision in bytes (BF16=2, INT8=1) + page_size_tokens: Tokens per page (default 64) + + Returns: + MLALayout: Layout instance for MLA architecture + + Raises: + ValueError: If model configuration is invalid + """ + required_fields = ['num_layers', 'kv_lora_rank', 'qk_rope_head_dim', + 'index_head_dim', 'kv_precision_bytes'] + for field in required_fields: + if field not in model_config: + raise ValueError(f"Missing required field: {field}") + + return MLALayout( + num_layers=model_config['num_layers'], + kv_lora_rank=model_config['kv_lora_rank'], + qk_rope_head_dim=model_config['qk_rope_head_dim'], + index_head_dim=model_config['index_head_dim'], + precision_bytes=model_config['kv_precision_bytes'], + page_size_tokens=page_size_tokens, + ) diff --git a/benchmarks/storage_benchmark_v1/storage/__init__.py b/benchmarks/storage_benchmark_v1/storage/__init__.py new file mode 100644 index 0000000000..40fe3856c9 --- /dev/null +++ b/benchmarks/storage_benchmark_v1/storage/__init__.py @@ -0,0 +1,13 @@ +""" +KVCache Storage Module + +Provides storage backend implementations for KVCache systems. +""" + +from .interface import Storage +from .disk import DiskHashTable + +__all__ = [ + 'Storage', + 'DiskHashTable', +] diff --git a/benchmarks/storage_benchmark_v1/storage/disk.py b/benchmarks/storage_benchmark_v1/storage/disk.py new file mode 100644 index 0000000000..b54de98a64 --- /dev/null +++ b/benchmarks/storage_benchmark_v1/storage/disk.py @@ -0,0 +1,333 @@ +""" +Simple SSD Hash Table Storage + +Each key maps to a complete page entry. +""" + +import os +import time +from pathlib import Path +from typing import Dict, Any + +from .interface import Storage + + +def calc_percentiles(data): + """Calculate latency percentiles""" + if not data: + return {'avg_ms': 0, 'p50_ms': 0, 'p95_ms': 0, 'p99_ms': 0} + import statistics + sorted_data = sorted(data) + + def get_percentile(p): + if len(sorted_data) == 1: + return sorted_data[0] + rank = (len(sorted_data) - 1) * (p / 100) + lower = int(rank) + upper = min(lower + 1, len(sorted_data) - 1) + weight = rank - lower + return sorted_data[lower] * (1.0 - weight) + sorted_data[upper] * weight + + return { + 'avg_ms': statistics.mean(data), + 'p50_ms': get_percentile(50), + 'p95_ms': get_percentile(95), + 'p99_ms': get_percentile(99), + } + + +class DiskHashTable(Storage): + """Simple disk-based hash table + + File Layout: + Page 0: offset = 0 * page_size + Page 1: offset = 1 * page_size + Page 2: offset = 2 * page_size + ... + Page N: offset = N * page_size + + When max_pages < actual page_id range, uses modulo mapping: + actual_page_id = page_id % max_pages + This allows simulating large traces with limited storage. + """ + + def __init__(self, storage_dir: str, page_size: int, + max_pages: int = 100000, + fsync_mode: str = 'batch', fsync_batch_size: int = 100): + """Initialize disk hash table + + Args: + storage_dir: Storage directory + page_size: Size of each entry (page) in bytes + max_pages: Maximum number of entries (creates circular mapping if trace is larger) + fsync_mode: When to fsync ('batch', 'always', 'end', 'none') + fsync_batch_size: Writes between fsync + """ + self.storage_dir = Path(storage_dir) + self.page_size = page_size + self.max_pages = max_pages + self.max_page_id = max_pages - 1 + self.fsync_mode = fsync_mode + self.fsync_batch_size = fsync_batch_size + self.storage_dir.mkdir(parents=True, exist_ok=True) + self.storage_file = self.storage_dir / "data.bin" + self._allocate_file() + self.fd = None + self._buffer = os.urandom(page_size) + self.stats = { + 'read_count': 0, + 'write_count': 0, + 'read_bytes': 0, + 'write_bytes': 0, + 'read_latencies_ms': [], + 'write_latencies_ms': [], + 'read_time_s': 0.0, + 'write_time_s': 0.0, + 'sync_count': 0, + 'hit': 0, + 'miss': 0, + } + self._pending_syncs = 0 + self._written_pages: set = set() + + def _allocate_file(self): + """Pre-allocate disk space efficiently""" + file_size = self.max_pages * self.page_size + if not self.storage_file.exists(): + print(f" [Storage] Creating file: {self.storage_file}") + print(f" [Storage] Requested size: {file_size / (1024**3):.2f} GB ({self.max_pages:,} pages × {self.page_size} bytes = {file_size:,} bytes)") + # Use fallocate for efficient preallocation (Linux) + fd = os.open(self.storage_file, os.O_RDWR | os.O_CREAT | os.O_TRUNC, 0o644) + try: + # Try fallocate first (Linux specific, much faster) + try: + import fcntl + fcntl.fallocate(fd, 0, file_size) + method = "fallocate" + except (ImportError, AttributeError, OSError): + # Fallback to seek+write method + os.lseek(fd, file_size - 1, os.SEEK_SET) + os.write(fd, b'\0') + os.fsync(fd) + method = "seek+write" + finally: + os.close(fd) + actual_size = self.storage_file.stat().st_size if self.storage_file.exists() else 0 + print(f" [Storage] Pre-allocated {actual_size / (1024**3):.2f} GB using {method}") + else: + actual_size = self.storage_file.stat().st_size + actual_pages = actual_size // self.page_size + print(f" [Storage] Reusing existing file: {self.storage_file}") + print(f" [Storage] Current file size: {actual_size / (1024**3):.2f} GB ({actual_pages:,} pages × {self.page_size} bytes = {actual_size:,} bytes)") + + def _get_fd(self): + if self.fd is None: + print(f" [Storage] Opening file: {self.storage_file}") + self.fd = os.open(self.storage_file, os.O_RDWR | os.O_CREAT, 0o644) + if self.storage_file.exists(): + actual_size = self.storage_file.stat().st_size + actual_pages = actual_size // self.page_size + print(f" [Storage] File size: {actual_size / (1024**3):.2f} GB ({actual_pages:,} pages × {self.page_size} bytes = {actual_size:,} bytes)") + return self.fd + + # ======================================================================== + # Core operations + # ======================================================================== + + def _map_page_id(self, page_id: int) -> int: + """Map logical page_id to physical page_id using modulo + + This allows simulating large traces with limited storage space. + + Args: + page_id: Logical page_id (from trace) + + Returns: + Physical page_id in storage (0 to max_pages-1) + """ + return page_id % self.max_pages + + def read(self, page_id: int, offset_in_page: int = 0, length: int = None) -> float: + """Read entry from disk + + Args: + page_id: Logical page ID (hash_id from trace) + offset_in_page: Offset within the page (default: 0) + length: Number of bytes to read (default: entire page) + + Returns: + Read latency in ms + """ + if length is None: + length = self.page_size - offset_in_page + + # Validate parameters + if offset_in_page < 0 or offset_in_page >= self.page_size: + raise ValueError(f"offset_in_page {offset_in_page} out of range [0, {self.page_size})") + if length <= 0 or offset_in_page + length > self.page_size: + raise ValueError(f"length {length} invalid with offset_in_page {offset_in_page} (page_size={self.page_size})") + + # Map to physical page_id and calculate offset + physical_page_id = self._map_page_id(page_id) + offset = physical_page_id * self.page_size + offset_in_page + start = time.perf_counter() + + try: + fd = self._get_fd() + os.pread(fd, length, offset) + + latency = (time.perf_counter() - start) * 1000.0 + self.stats['read_count'] += 1 + self.stats['read_bytes'] += length + self.stats['read_latencies_ms'].append(latency) + self.stats['read_time_s'] += latency / 1000.0 + self.stats['hit'] += 1 + return latency + except OSError as e: + print(f"Read error (page_id={page_id}, physical_page_id={physical_page_id}): {e}") + return 0.0 + + def write(self, page_id: int, offset_in_page: int = 0, length: int = None) -> float: + """Write entry to disk + + Args: + page_id: Logical page ID (hash_id from trace) + offset_in_page: Offset within the page (default: 0) + length: Number of bytes to write (default: entire page) + + Returns: + Write latency in ms + """ + if length is None: + length = self.page_size - offset_in_page + + # Validate parameters + if offset_in_page < 0 or offset_in_page >= self.page_size: + raise ValueError(f"offset_in_page {offset_in_page} out of range [0, {self.page_size})") + if length <= 0 or offset_in_page + length > self.page_size: + raise ValueError(f"length {length} invalid with offset_in_page {offset_in_page} (page_size={self.page_size})") + + # Map to physical page_id and calculate offset + physical_page_id = self._map_page_id(page_id) + offset = physical_page_id * self.page_size + offset_in_page + start = time.perf_counter() + + try: + fd = self._get_fd() + # Use corresponding portion of buffer + os.pwrite(fd, self._buffer[:length], offset) + write_done = time.perf_counter() + + # Fsync + if self.fsync_mode == 'always': + os.fsync(fd) + self.stats['sync_count'] += 1 + self._pending_syncs = 0 + latency = (time.perf_counter() - start) * 1000.0 + elif self.fsync_mode == 'batch': + self._pending_syncs += 1 + if self._pending_syncs >= self.fsync_batch_size: + os.fsync(fd) + self.stats['sync_count'] += 1 + self._pending_syncs = 0 + latency = (write_done - start) * 1000.0 + else: + latency = (write_done - start) * 1000.0 + + self._written_pages.add(page_id) + self.stats['write_count'] += 1 + self.stats['write_bytes'] += length + self.stats['write_latencies_ms'].append(latency) + self.stats['write_time_s'] += latency / 1000.0 + self.stats['miss'] += 1 + return latency + except OSError as e: + print(f"Write error (page_id={page_id}, offset_in_page={offset_in_page}, length={length}): {e}") + return 0.0 + + # ======================================================================== + # Storage interface methods + # ======================================================================== + + def exists(self, page_id: int) -> bool: + """Check if logical page has been written + + Args: + page_id: Logical page ID (hash_id from trace) + + Returns: + True if this logical page has been written before + """ + return page_id in self._written_pages + + def delete(self, page_id: int) -> bool: + """Delete page (no-op in direct mapping) + + Args: + page_id: Page ID to delete + + Returns: + True (always succeeds in direct mapping) + """ + # No-op since we don't track which pages have been written + return True + + def get_stats(self) -> Dict[str, Any]: + """Get statistics""" + import statistics + + def calc_stats(latencies): + 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, + 'time_s': self.stats['read_time_s'], + **calc_stats(self.stats['read_latencies_ms']) + }, + 'write': { + 'count': self.stats['write_count'], + 'mb': self.stats['write_bytes'] / 1024 / 1024, + 'time_s': self.stats['write_time_s'], + **calc_stats(self.stats['write_latencies_ms']) + }, + 'sync_count': self.stats['sync_count'], + 'max_pages': self.max_pages, + 'written_pages': len(self._written_pages), + 'page_hits': self.stats['hit'], + 'page_misses': self.stats['miss'], + } + + # ======================================================================== + # Resource management + # ======================================================================== + + def close(self, force_sync: bool = True): + """Close file""" + if force_sync and self.fsync_mode in ['end', 'batch']: + if self.fd is not None: + try: + os.fsync(self.fd) + self.stats['sync_count'] += 1 + except OSError: + pass + + if self.fd is not None: + try: + os.close(self.fd) + except OSError: + pass + self.fd = None + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() + return False diff --git a/benchmarks/storage_benchmark_v1/storage/interface.py b/benchmarks/storage_benchmark_v1/storage/interface.py new file mode 100644 index 0000000000..e7c2aa022a --- /dev/null +++ b/benchmarks/storage_benchmark_v1/storage/interface.py @@ -0,0 +1,93 @@ +""" +KVCache Storage Interface + +Simplified storage interface for KVCache benchmark. +Key = page_id (int), Value = fixed-size bytes. +""" + +from abc import ABC, abstractmethod +from typing import Dict, Any + + +class Storage(ABC): + """Abstract base class for KVCache storage + + Simplified design: + - Key: page_id (int, hash_id from trace) + - Value: fixed-size bytes (page_size) + - Direct mapping: page_id -> offset -> payload + + Core operations: read, write, exists, delete. + """ + + @abstractmethod + def read(self, page_id: int, offset_in_page: int = 0, length: int = None) -> float: + """Read page from disk + + Args: + page_id: Page ID (hash_id from trace) + offset_in_page: Offset within the page (default: 0) + length: Number of bytes to read (default: entire page) + + Returns: + Read latency in milliseconds + """ + pass + + @abstractmethod + def write(self, page_id: int, offset_in_page: int = 0, length: int = None) -> float: + """Write page to disk + + Args: + page_id: Page ID (hash_id from trace) + offset_in_page: Offset within the page (default: 0) + length: Number of bytes to write (default: entire page) + + Returns: + Write latency in milliseconds + """ + pass + + @abstractmethod + def exists(self, page_id: int) -> bool: + """Check if page exists + + Args: + page_id: Page ID (hash_id from trace) + + Returns: + True if page_id is within valid range + """ + pass + + @abstractmethod + def delete(self, page_id: int) -> bool: + """Delete page (no-op in direct mapping) + + Args: + page_id: Page ID to delete + + Returns: + True (always succeeds in direct mapping) + """ + pass + + @abstractmethod + def get_stats(self) -> Dict[str, Any]: + """Get storage statistics + + Returns: + Dictionary containing storage statistics + """ + pass + + def close(self): + """Close the storage and release resources""" + pass + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() + return False diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake new file mode 100644 index 0000000000..a9ccc3e787 --- /dev/null +++ b/cmake/Dependencies.cmake @@ -0,0 +1,74 @@ +# Build entry points own these targets and call each provider at most once. +# Consumer directories only link the resulting Mooncake::* targets. + +function(mooncake_provide_zstd) + find_path(MOONCAKE_ZSTD_INCLUDE_DIR zstd.h) + find_library(MOONCAKE_ZSTD_LIBRARY zstd) + if(NOT MOONCAKE_ZSTD_INCLUDE_DIR OR NOT MOONCAKE_ZSTD_LIBRARY) + message(FATAL_ERROR "zstd development files not found") + endif() + add_library(Mooncake::zstd UNKNOWN IMPORTED) + set_target_properties( + Mooncake::zstd + PROPERTIES IMPORTED_LOCATION "${MOONCAKE_ZSTD_LIBRARY}" + INTERFACE_INCLUDE_DIRECTORIES "${MOONCAKE_ZSTD_INCLUDE_DIR}") +endfunction() + +function(mooncake_provide_xxhash) + find_path(MOONCAKE_XXHASH_INCLUDE_DIR xxhash.h) + find_library(MOONCAKE_XXHASH_LIBRARY NAMES xxhash libxxhash) + if(NOT MOONCAKE_XXHASH_INCLUDE_DIR OR NOT MOONCAKE_XXHASH_LIBRARY) + message(FATAL_ERROR "xxHash development files not found") + endif() + add_library(Mooncake::xxhash UNKNOWN IMPORTED) + set_target_properties( + Mooncake::xxhash + PROPERTIES IMPORTED_LOCATION "${MOONCAKE_XXHASH_LIBRARY}" + INTERFACE_INCLUDE_DIRECTORIES "${MOONCAKE_XXHASH_INCLUDE_DIR}") +endfunction() + +function(mooncake_provide_liburing) + find_path(MOONCAKE_LIBURING_INCLUDE_DIR liburing.h) + find_library(MOONCAKE_LIBURING_LIBRARY uring) + if(MOONCAKE_LIBURING_INCLUDE_DIR AND MOONCAKE_LIBURING_LIBRARY) + add_library(Mooncake::liburing UNKNOWN IMPORTED) + set_target_properties( + Mooncake::liburing + PROPERTIES IMPORTED_LOCATION "${MOONCAKE_LIBURING_LIBRARY}" + INTERFACE_INCLUDE_DIRECTORIES + "${MOONCAKE_LIBURING_INCLUDE_DIR}") + endif() +endfunction() + +function(mooncake_provide_libzmq) + find_path(MOONCAKE_LIBZMQ_INCLUDE_DIR zmq.h) + find_library(MOONCAKE_LIBZMQ_LIBRARY NAMES zmq libzmq) + if(NOT MOONCAKE_LIBZMQ_INCLUDE_DIR OR NOT MOONCAKE_LIBZMQ_LIBRARY) + message(FATAL_ERROR "libzmq development files not found") + endif() + add_library(Mooncake::libzmq UNKNOWN IMPORTED) + set_target_properties( + Mooncake::libzmq + PROPERTIES IMPORTED_LOCATION "${MOONCAKE_LIBZMQ_LIBRARY}" + INTERFACE_INCLUDE_DIRECTORIES "${MOONCAKE_LIBZMQ_INCLUDE_DIR}") +endfunction() + +function(mooncake_provide_hiredis) + cmake_parse_arguments(ARG "REQUIRED" "" "" ${ARGN}) + find_path(MOONCAKE_HIREDIS_INCLUDE_DIR hiredis/hiredis.h) + find_library(MOONCAKE_HIREDIS_LIBRARY hiredis) + if(ARG_REQUIRED) + if(NOT MOONCAKE_HIREDIS_INCLUDE_DIR OR NOT MOONCAKE_HIREDIS_LIBRARY) + message(FATAL_ERROR "hiredis development files not found") + endif() + endif() + + if(MOONCAKE_HIREDIS_INCLUDE_DIR AND MOONCAKE_HIREDIS_LIBRARY) + add_library(Mooncake::hiredis UNKNOWN IMPORTED) + set_target_properties( + Mooncake::hiredis + PROPERTIES IMPORTED_LOCATION "${MOONCAKE_HIREDIS_LIBRARY}" + INTERFACE_INCLUDE_DIRECTORIES + "${MOONCAKE_HIREDIS_INCLUDE_DIR}") + endif() +endfunction() diff --git a/dependencies.sh b/dependencies.sh index 0c97cf2730..79fb6084de 100755 --- a/dependencies.sh +++ b/dependencies.sh @@ -23,8 +23,8 @@ NC="\033[0m" # No Color # Configuration REPO_ROOT=`pwd` GITHUB_PROXY=${GITHUB_PROXY:-"https://github.com"} -GOVER=1.23.8 -YALANTINGLIBS_VERSION=0.5.7 +GOVER=1.25.9 +OS_RELEASE_FILE=${OS_RELEASE_FILE:-/etc/os-release} # Function to print section headers print_section() { @@ -49,23 +49,57 @@ check_success() { fi } +read_os_release_value() { + local key="$1" + awk -F= -v key="$key" ' + $1 == key { + value = $0 + sub(/^[^=]*=/, "", value) + gsub(/^"|"$/, "", value) + print value + exit + } + ' "$OS_RELEASE_FILE" +} + +# Function to detect OS +detect_os() { + if [ -f "$OS_RELEASE_FILE" ]; then + ID=$(read_os_release_value ID) + VERSION_ID=$(read_os_release_value VERSION_ID) + OS=$(echo "$ID" | tr '[:upper:]' '[:lower:]') + OS_VERSION=$VERSION_ID + elif [ -f /etc/redhat-release ]; then + OS="centos" + else + print_error "Cannot detect OS. Supported OS: Ubuntu, Debian, CentOS, RHEL, Rocky, AlmaLinux, EulerOS, and openEuler." + fi + + echo -e "${GREEN}Detected OS: $OS ${OS_VERSION:-unknown}${NC}" +} + if [ $(id -u) -ne 0 ]; then print_error "Require root permission, try sudo ./dependencies.sh" fi # Parse command line arguments SKIP_CONFIRM=false +INSTALL_SPDK=false for arg in "$@"; do case $arg in -y|--yes) SKIP_CONFIRM=true ;; + --with-spdk) + INSTALL_SPDK=true + ;; -h|--help) echo -e "${YELLOW}Mooncake Dependencies Installer${NC}" echo -e "Usage: ./dependencies.sh [OPTIONS]" echo -e "\nOptions:" - echo -e " -y, --yes Skip confirmation and install all dependencies" - echo -e " -h, --help Show this help message and exit" + echo -e " -y, --yes Skip confirmation and install all dependencies" + echo -e " --with-spdk Install SPDK for NVMe-oF support" + echo -e " -h, --help Show this help message and exit" exit 0 ;; esac @@ -76,9 +110,11 @@ 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 " - yalantinglibs" -echo -e " - Git submodules" +echo -e " - Git submodules (pybind11)" echo -e " - Go $GOVER" +if [ "$INSTALL_SPDK" = true ]; then + echo -e " - SPDK (for NVMe-oF support)" +fi echo # Ask for confirmation unless -y flag is used @@ -91,131 +127,118 @@ if [ "$SKIP_CONFIRM" = false ]; then fi fi +# Detect OS +detect_os # Update package lists print_section "Updating package lists" -apt-get update -check_success "Failed to update package lists" +if [ "$OS" = "ubuntu" ] || [ "$OS" = "debian" ]; then + apt-get update + check_success "Failed to update package lists" +elif [ "$OS" = "centos" ] || [ "$OS" = "rhel" ] || [ "$OS" = "rocky" ] || [ "$OS" = "almalinux" ] || [ "$OS" = "euleros" ] || [ "$OS" = "openeuler" ]; then + yum install -y dnf-plugins-core epel-release || true + yum config-manager --set-enabled powertools || yum config-manager --set-enabled crb || true + yum clean all + yum makecache + check_success "Failed to update package lists" +else + print_error "Unsupported OS: $OS" +fi # Install system packages print_section "Installing system packages" echo -e "${YELLOW}This may take a few minutes...${NC}" -SYSTEM_PACKAGES="build-essential \ - cmake \ - git \ - wget \ - unzip \ - libibverbs-dev \ - libgoogle-glog-dev \ - libgtest-dev \ - libjsoncpp-dev \ - libunwind-dev \ - libnuma-dev \ - libpython3-dev \ - libboost-all-dev \ - libssl-dev \ - libgrpc-dev \ - libgrpc++-dev \ - libprotobuf-dev \ - libyaml-cpp-dev \ - protobuf-compiler-grpc \ - libcurl4-openssl-dev \ - libhiredis-dev \ - liburing-dev \ - libjemalloc-dev \ - libmsgpack-dev \ - libzstd-dev \ - libasio-dev \ - pkg-config \ - patchelf \ - libc6-dev \ - libc-bin" - -apt-get install -y $SYSTEM_PACKAGES -check_success "Failed to install system packages" -print_success "System packages installed successfully" - -# Install yalantinglibs -print_section "Installing yalantinglibs" - -# Check if thirdparties directory exists -if [ ! -d "${REPO_ROOT}/thirdparties" ]; then - mkdir -p "${REPO_ROOT}/thirdparties" - check_success "Failed to create thirdparties directory" -fi - -# 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" +if [ "$OS" = "ubuntu" ] || [ "$OS" = "debian" ]; then + SYSTEM_PACKAGES="build-essential \ + cmake \ + ninja-build \ + git \ + wget \ + unzip \ + libibverbs-dev \ + libgoogle-glog-dev \ + libjsoncpp-dev \ + libunwind-dev \ + libnuma-dev \ + libpython3-dev \ + libboost-all-dev \ + libssl-dev \ + libgrpc-dev \ + libgrpc++-dev \ + libprotobuf-dev \ + libyaml-cpp-dev \ + protobuf-compiler-grpc \ + libcurl4-openssl-dev \ + libhiredis-dev \ + liburing-dev \ + libjemalloc-dev \ + libmsgpack-dev \ + libzmq3-dev \ + libzstd-dev \ + libasio-dev \ + libxxhash-dev \ + pkg-config \ + patchelf \ + libc6-dev \ + libc-bin" + + apt-get install -y $SYSTEM_PACKAGES + check_success "Failed to install system packages" + +elif [ "$OS" = "centos" ] || [ "$OS" = "rhel" ] || [ "$OS" = "rocky" ] || [ "$OS" = "almalinux" ] || [ "$OS" = "euleros" ] || [ "$OS" = "openeuler" ]; then + SYSTEM_PACKAGES="@development \ + cmake \ + ninja-build \ + git \ + wget \ + rdma-core-devel \ + glog-devel \ + gflags-devel \ + jsoncpp-devel \ + libunwind-devel \ + numactl-devel \ + python3-devel \ + boost1.78-devel \ + openssl-devel \ + protobuf-devel \ + yaml-cpp-devel \ + libcurl-devel \ + hiredis-devel \ + liburing-devel \ + jemalloc-devel \ + msgpack-devel \ + libzstd-devel \ + pkgconf-pkg-config \ + elfutils-libelf-devel \ + patchelf \ + xxhash-devel \ + libbsd-devel" + + yum install -y $SYSTEM_PACKAGES + check_success "Failed to install system packages" +else + print_error "Unsupported OS: $OS" 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" - -echo "Configuring yalantinglibs..." -cmake .. -DBUILD_EXAMPLES=OFF -DBUILD_BENCHMARK=OFF -DBUILD_UNIT_TESTS=OFF -check_success "Failed to configure yalantinglibs" - -echo "Building yalantinglibs (using $(nproc) cores)..." -cmake --build . -j$(nproc) -check_success "Failed to build yalantinglibs" - -echo "Installing yalantinglibs..." -cmake --install . -check_success "Failed to install yalantinglibs" - -print_success "yalantinglibs installed successfully" +print_success "System packages 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" + 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" - fi + print_success "Git submodules initialized and updated successfully" else echo -e "${YELLOW}No .gitmodules file found. Skipping...${NC}" exit 1 @@ -224,18 +247,21 @@ fi print_section "Verifying essential build tools" # Verify getconf and ldd (required for glibc version detection in build_wheel.sh) -# Both are provided by libc-bin, which is included in SYSTEM_PACKAGES -if ! command -v getconf >/dev/null 2>&1; then - print_error "getconf not found after installing system packages. This should not happen." -fi -if ! command -v ldd >/dev/null 2>&1; then - print_error "ldd not found after installing system packages. This should not happen." +if [ "$OS" = "ubuntu" ] || [ "$OS" = "debian" ]; then + if ! command -v getconf >/dev/null 2>&1; then + print_error "getconf not found after installing system packages. This should not happen." + fi + if ! command -v ldd >/dev/null 2>&1; then + print_error "ldd not found after installing system packages. This should not happen." + fi + print_success "getconf found: $(getconf --version 2>&1 | head -1)" + print_success "ldd found: $(ldd --version 2>&1 | head -1)" fi -print_success "getconf found: $(getconf --version 2>&1 | head -1)" -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 @@ -246,24 +272,46 @@ install_go() { echo "Unsupported architecture: $ARCH" exit 1 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 + 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 [[ "$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 + echo "Installing Go $GOVER..." - tar -C /usr/local -xzf go$GOVER.linux-$ARCH.tar.gz + tar -C /usr/local -xzf "${GO_TARBALL}" check_success "Failed to install Go $GOVER" - # Clean up downloaded file - rm -f go$GOVER.linux-$ARCH.tar.gz + rm -f "${GO_TARBALL}" check_success "Failed to clean up Go installation file" print_success "Go $GOVER installed successfully" } -# Check if Go is already installed if command -v go &> /dev/null; then GO_VERSION=$(go version | awk '{print $3}') if [[ "$GO_VERSION" == "go$GOVER" ]]; then @@ -283,6 +331,82 @@ 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 +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 + +# Install SPDK if requested +if [ "$INSTALL_SPDK" = true ]; then + print_section "Installing SPDK" + + cd "${REPO_ROOT}/extern" + check_success "Failed to change to extern directory" + + # Remove existing SPDK if present + if [ -d "spdk" ]; then + echo -e "${YELLOW}SPDK directory already exists. Removing for fresh install...${NC}" + rm -rf spdk + check_success "Failed to remove existing SPDK directory" + fi + + # Clone SPDK + echo "Cloning SPDK from ${GITHUB_PROXY}/spdk/spdk.git..." + git clone ${GITHUB_PROXY}/spdk/spdk.git + check_success "Failed to clone SPDK" + + cd spdk + check_success "Failed to change to SPDK directory" + + # Checkout specific version + echo "Checking out SPDK version v23.01.1..." + git checkout v23.01.1 + check_success "Failed to checkout SPDK version v23.01.1" + + # Initialize submodules + echo "Initializing SPDK submodules..." + git submodule update --init + check_success "Failed to initialize SPDK submodules" + + # Install SPDK dependencies + echo "Installing SPDK dependencies..." + ./scripts/pkgdep.sh + check_success "Failed to install SPDK dependencies" + + # Configure SPDK with RDMA support + echo "Configuring SPDK with RDMA support..." + ./configure --with-rdma + check_success "Failed to configure SPDK" + + # Build SPDK + echo "Building SPDK (using $(nproc) cores)..." + make -j$(nproc) + check_success "Failed to build SPDK" + + # Install SPDK + echo "Installing SPDK..." + make install + check_success "Failed to install SPDK" + + # Copy DPDK libraries to system library path + if ls dpdk/build/lib/*.a >/dev/null 2>&1; then + echo "Copying DPDK libraries to /usr/local/lib..." + cp dpdk/build/lib/*.a /usr/local/lib/ + check_success "Failed to copy DPDK libraries" + fi + + print_success "SPDK installed successfully" + cd "${REPO_ROOT}" +fi + # Return to the repository root cd "${REPO_ROOT}" @@ -291,9 +415,15 @@ print_section "Installation Complete" echo -e "${GREEN}All dependencies have been successfully installed!${NC}" echo -e "The following components were installed:" echo -e " ${GREEN}✓${NC} System packages" -echo -e " ${GREEN}✓${NC} yalantinglibs" echo -e " ${GREEN}✓${NC} Git submodules" echo -e " ${GREEN}✓${NC} Go $GOVER" +if [ "$INSTALL_SPDK" = true ]; then + echo -e " ${GREEN}✓${NC} SPDK (v23.01.1)" +fi echo echo -e "You can now build and run Mooncake." echo -e "${YELLOW}Note: You may need to restart your terminal or run 'source ~/.bashrc' to use Go.${NC}" + +if [ "$INSTALL_SPDK" = true ]; then + echo -e "${YELLOW}Note: SPDK requires hugepages and RDMA configuration. Please refer to SPDK documentation for setup.${NC}" +fi diff --git a/docker/master-cuda13.Dockerfile b/docker/master-cuda13.Dockerfile new file mode 100644 index 0000000000..c01903fbc7 --- /dev/null +++ b/docker/master-cuda13.Dockerfile @@ -0,0 +1,55 @@ +# syntax=docker/dockerfile:1 + +# Keep docker/master.Dockerfile and docker/master-cuda13.Dockerfile in sync: they must +# differ ONLY in the three CUDA-flavor lines (the cudalibs FROM, the libcudart COPY and +# the pip package name). The publish workflow diffs them with those lines normalized. + +# Stage cudalibs: take stub libcuda and libcudart from the CUDA devel image +FROM nvidia/cuda:13.0.3-devel-ubuntu22.04@sha256:3869b846a8cc495ce11c172d87cfc0da8874b910d14a9810bec6b6182e9ee9f8 AS cudalibs + +# Final image. Must be trixie (glibc 2.41): the aarch64 wheel is manylinux_2_39 +# (needs glibc >= 2.39), which bookworm (2.36) cannot satisfy. +FROM python:3.12-slim-trixie@sha256:d764629ce0ddd8c71fd371e9901efb324a95789d2315a47db7e4d27e78f1b0e9 + +# Build args: mooncake version, pip index URL +ARG MOONCAKE_VERSION +ARG PIP_INDEX_URL=https://pypi.org/simple + +# Install runtime system libraries and tini. +# ibverbs-providers ships /usr/lib//libmlx5.so.1 plus the libibverbs provider +# plugins, and both are required: +# - at load time, because since 0.3.12 engine.so / store.so / mooncake_master carry a +# hard DT_NEEDED on libmlx5.so.1 (the transport links mlx5 for the IBGDA / mlx5dv +# DevX path) and auditwheel deliberately does not vendor RDMA libraries into the +# wheel -- without it `import mooncake.engine` fails outright, even for TCP-only use; +# - at run time, because libibverbs claims devices through those provider plugins: +# without the package ibv_get_device_list() returns 0 devices even when the mlx5 +# devices are visible in /sys/class/infiniband and /dev/infiniband is passed in. +# It must come from apt next to libibverbs1: both are built from the rdma-core source +# package and share a private provider ABI, so their versions have to match. +RUN apt-get update && apt-get install -y --no-install-recommends \ + ca-certificates libibverbs1 ibverbs-providers libnuma1 libcurl4t64 libstdc++6 tini \ + && rm -rf /var/lib/apt/lists/* + +# Copy stub libcuda and libcudart into the loader's default path, refresh the link cache +COPY --from=cudalibs /usr/local/cuda/lib64/stubs/libcuda.so /usr/local/lib/libcuda.so.1 +COPY --from=cudalibs /usr/local/cuda/lib64/libcudart.so.13 /usr/local/lib/libcudart.so.13 +RUN ldconfig + +# Install mooncake, remove torch EP/PG extensions (ep_*/pg_*), chown the package dir to uid 65532 (all in one layer) +RUN pip install --no-cache-dir --index-url "${PIP_INDEX_URL}" \ + mooncake-transfer-engine-cuda13==${MOONCAKE_VERSION} \ + && PKG="$(python3 -c 'import mooncake,os;print(os.path.dirname(mooncake.__file__))')" \ + && rm -f "$PKG"/ep_*.so "$PKG"/pg_*.so \ + && chown -R 65532:65532 "$PKG" + +# Create a HOME owned by uid 65532 and set it as WORKDIR +ENV HOME=/home/nonroot +RUN mkdir -p /home/nonroot && chown 65532:65532 /home/nonroot +WORKDIR /home/nonroot + +USER 65532:65532 + +# tini as PID 1 to forward signals; default into bash +ENTRYPOINT ["tini", "-g", "--"] +CMD ["bash"] diff --git a/docker/master.Dockerfile b/docker/master.Dockerfile new file mode 100644 index 0000000000..39b7e163c7 --- /dev/null +++ b/docker/master.Dockerfile @@ -0,0 +1,55 @@ +# syntax=docker/dockerfile:1 + +# Keep docker/master.Dockerfile and docker/master-cuda13.Dockerfile in sync: they must +# differ ONLY in the three CUDA-flavor lines (the cudalibs FROM, the libcudart COPY and +# the pip package name). The publish workflow diffs them with those lines normalized. + +# Stage cudalibs: take stub libcuda and libcudart from the CUDA devel image +FROM nvidia/cuda:12.8.1-devel-ubuntu22.04@sha256:a99a1860ba8e2916e5c3e73b72ec4c4301653a84586e05bfc9a2aa2d58027e97 AS cudalibs + +# Final image. Must be trixie (glibc 2.41): the aarch64 wheel is manylinux_2_39 +# (needs glibc >= 2.39), which bookworm (2.36) cannot satisfy. +FROM python:3.12-slim-trixie@sha256:d764629ce0ddd8c71fd371e9901efb324a95789d2315a47db7e4d27e78f1b0e9 + +# Build args: mooncake version, pip index URL +ARG MOONCAKE_VERSION +ARG PIP_INDEX_URL=https://pypi.org/simple + +# Install runtime system libraries and tini. +# ibverbs-providers ships /usr/lib//libmlx5.so.1 plus the libibverbs provider +# plugins, and both are required: +# - at load time, because since 0.3.12 engine.so / store.so / mooncake_master carry a +# hard DT_NEEDED on libmlx5.so.1 (the transport links mlx5 for the IBGDA / mlx5dv +# DevX path) and auditwheel deliberately does not vendor RDMA libraries into the +# wheel -- without it `import mooncake.engine` fails outright, even for TCP-only use; +# - at run time, because libibverbs claims devices through those provider plugins: +# without the package ibv_get_device_list() returns 0 devices even when the mlx5 +# devices are visible in /sys/class/infiniband and /dev/infiniband is passed in. +# It must come from apt next to libibverbs1: both are built from the rdma-core source +# package and share a private provider ABI, so their versions have to match. +RUN apt-get update && apt-get install -y --no-install-recommends \ + ca-certificates libibverbs1 ibverbs-providers libnuma1 libcurl4t64 libstdc++6 tini \ + && rm -rf /var/lib/apt/lists/* + +# Copy stub libcuda and libcudart into the loader's default path, refresh the link cache +COPY --from=cudalibs /usr/local/cuda/lib64/stubs/libcuda.so /usr/local/lib/libcuda.so.1 +COPY --from=cudalibs /usr/local/cuda/lib64/libcudart.so.12 /usr/local/lib/libcudart.so.12 +RUN ldconfig + +# Install mooncake, remove torch EP/PG extensions (ep_*/pg_*), chown the package dir to uid 65532 (all in one layer) +RUN pip install --no-cache-dir --index-url "${PIP_INDEX_URL}" \ + mooncake-transfer-engine==${MOONCAKE_VERSION} \ + && PKG="$(python3 -c 'import mooncake,os;print(os.path.dirname(mooncake.__file__))')" \ + && rm -f "$PKG"/ep_*.so "$PKG"/pg_*.so \ + && chown -R 65532:65532 "$PKG" + +# Create a HOME owned by uid 65532 and set it as WORKDIR +ENV HOME=/home/nonroot +RUN mkdir -p /home/nonroot && chown 65532:65532 /home/nonroot +WORKDIR /home/nonroot + +USER 65532:65532 + +# tini as PID 1 to forward signals; default into bash +ENTRYPOINT ["tini", "-g", "--"] +CMD ["bash"] diff --git a/docker/mooncake.Dockerfile b/docker/mooncake.Dockerfile index 33542086fb..9ade9603d1 100644 --- a/docker/mooncake.Dockerfile +++ b/docker/mooncake.Dockerfile @@ -9,62 +9,77 @@ ARG UBUNTU_VERSION=22.04 FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu${UBUNTU_VERSION} AS builder ENV DEBIAN_FRONTEND=noninteractive \ - PYTHONUNBUFFERED=1 + PYTHONUNBUFFERED=1 \ + PIP_NO_CACHE_DIR=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" +ARG EP_TORCH_VERSIONS="2.13.0" +ARG TORCH_CUDA_ARCH_LIST="" +# CI can opt in to removing /workspace/build from the builder layer. +ARG CLEAN_BUILD_ARTIFACTS=0 ENV PYTHON_VERSION=${PYTHON_VERSION} \ - BUILD_WITH_EP=1 \ EP_TORCH_VERSIONS=${EP_TORCH_VERSIONS} \ TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} \ PATH="/usr/local/go/bin:${PATH}" -# Install base build utilities and python bindings +# Install base build utilities and the requested Python version via deadsnakes PPA RUN apt-get update && \ apt-get install -y --no-install-recommends \ ca-certificates \ curl \ git \ - python3 \ - python3-dev \ - python3-pip \ - python-is-python3 \ + ninja-build \ + software-properties-common \ 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 COPY . /workspace -# Install Mooncake dependencies (yalantinglibs, Go, etc.) +# Install Mooncake dependencies (submodules, Go, etc.) RUN bash dependencies.sh -y -# Configure & build Mooncake +# Configure and build the wheel in one layer, then remove build/ only after +# auditwheel has resolved libraries from it. The large intermediate tree is +# therefore not retained in the builder image or BuildKit cache. RUN mkdir -p build && \ cd build && \ - cmake .. \ + cmake -G Ninja .. \ -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 . -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 && \ + cmake --build . && \ + cd /workspace && \ + export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH && \ export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH && \ export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH && \ mkdir -p build/mooncake-transfer-engine/nvlink-allocator && \ cd mooncake-transfer-engine/nvlink-allocator && \ - bash build.sh ../../build/mooncake-transfer-engine/nvlink-allocator/ - -# Build the Python wheel from local sources -RUN OUTPUT_DIR=dist ./scripts/build_wheel.sh + bash build.sh ../../build/mooncake-transfer-engine/nvlink-allocator/ && \ + cd /workspace && \ + OUTPUT_DIR=dist ./scripts/build_wheel.sh && \ + if [ "${CLEAN_BUILD_ARTIFACTS}" = "1" ]; then \ + rm -rf build; \ + fi ############################################################################### # Stage 2: install the freshly built wheel into a runtime image @@ -75,11 +90,17 @@ ENV DEBIAN_FRONTEND=noninteractive \ PYTHONUNBUFFERED=1 \ PIP_NO_CACHE_DIR=1 -# Install runtime dependencies required by Mooncake +# 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 RUN apt-get update && \ apt-get install -y --no-install-recommends \ - python3 \ - python3-pip \ + ca-certificates \ + curl \ + software-properties-common \ ibverbs-providers \ rdma-core \ libibverbs1 \ @@ -88,10 +109,19 @@ 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 python3 -m pip install --no-cache-dir /tmp/mooncake-wheel/*.whl && rm -rf /tmp/mooncake-wheel /root/.cache/pip +COPY --chmod=755 scripts/check_hicache_hugepage_requirements.py /usr/local/bin/mooncake-hicache-sizing +RUN python${PYTHON_VERSION} -m pip install --no-cache-dir /tmp/mooncake-wheel/*.whl && rm -rf /tmp/mooncake-wheel /root/.cache/pip CMD ["/bin/bash"] diff --git a/docker/musa.Dockerfile b/docker/musa.Dockerfile index d47564a321..1aff6ce363 100644 --- a/docker/musa.Dockerfile +++ b/docker/musa.Dockerfile @@ -3,48 +3,49 @@ ############################################################################### # Stage 1: build Mooncake from source and produce a Python wheel ############################################################################### -ARG MUSA_VERSION=rc4.3.0 +ARG TORCH_VERSION=2.9.1.post1 +ARG PYTHON_VERSION=3.10 +ARG MUSA_VERSION=5.2.0 +ARG MUSA_ARCH=mp31 ARG UBUNTU_VERSION=22.04 +ARG BASE_IMAGE=registry.mthreads.com/mcconline/inference/pytorch:${TORCH_VERSION}-py${PYTHON_VERSION}-musa${MUSA_VERSION}-${MUSA_ARCH}-devel-ubuntu${UBUNTU_VERSION}-amd64 -FROM mthreads/musa:${MUSA_VERSION}-devel-ubuntu${UBUNTU_VERSION}-amd64 AS builder +FROM ${BASE_IMAGE} AS builder ENV DEBIAN_FRONTEND=noninteractive \ PYTHONUNBUFFERED=1 -ARG PYTHON_VERSION=3.10 +ARG PYTHON_VERSION ARG CMAKE_BUILD_TYPE=Release +ARG PYPI_INDEX_URL=https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple +# Empty, unlike the CUDA image: the EP/PG extensions must build against the +# MUSA torch already installed here. Any value makes BuildEpExt.cmake +# pip-install that version from download.pytorch.org, replacing the MUSA torch +# stack with a CUDA build. +ARG EP_TORCH_VERSIONS="" ENV PYTHON_VERSION=${PYTHON_VERSION} \ + EP_TORCH_VERSIONS=${EP_TORCH_VERSIONS} \ + PIP_INDEX_URL=${PYPI_INDEX_URL} \ PATH="/usr/local/go/bin:${PATH}" -# Install base build utilities and python bindings -RUN apt-get update && \ - apt-get install -y --no-install-recommends \ - ca-certificates \ - curl \ - git \ - python3 \ - python3-dev \ - python3-pip \ - python-is-python3 \ - pkg-config && \ - rm -rf /var/lib/apt/lists/* - WORKDIR /workspace COPY . /workspace -# Install Mooncake dependencies (yalantinglibs, Go, etc.) +# Install Mooncake dependencies (submodules, Go, etc.) RUN bash dependencies.sh -y # Configure & build Mooncake RUN mkdir -p build && \ cd build && \ - cmake .. \ + cmake -G Ninja .. \ -DBUILD_UNIT_TESTS=OFF \ -DUSE_HTTP=ON \ -DUSE_ETCD=ON \ -DUSE_MUSA=ON \ + -DWITH_EP=ON \ -DSTORE_USE_ETCD=ON \ + -DPython3_EXECUTABLE=/usr/bin/python${PYTHON_VERSION} \ -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE} && \ cmake --build . -j"$(nproc)" @@ -59,17 +60,18 @@ RUN OUTPUT_DIR=dist ./scripts/build_wheel.sh ############################################################################### # Stage 2: install the freshly built wheel into a runtime image ############################################################################### -FROM mthreads/musa:${MUSA_VERSION}-devel-ubuntu${UBUNTU_VERSION}-amd64 AS runtime +FROM ${BASE_IMAGE} AS runtime ENV DEBIAN_FRONTEND=noninteractive \ PYTHONUNBUFFERED=1 \ PIP_NO_CACHE_DIR=1 +ARG PYTHON_VERSION +ARG PYPI_INDEX_URL=https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple + # Install runtime dependencies required by Mooncake RUN apt-get update && \ apt-get install -y --no-install-recommends \ - python3 \ - python3-pip \ ibverbs-providers \ rdma-core \ libibverbs1 \ @@ -82,6 +84,6 @@ RUN apt-get update && \ # Copy wheels produced in builder stage and install them via pip COPY --from=builder /workspace/mooncake-wheel/dist /tmp/mooncake-wheel -RUN python3 -m pip install --no-cache-dir /tmp/mooncake-wheel/*.whl && rm -rf /tmp/mooncake-wheel /root/.cache/pip +RUN PIP_INDEX_URL=${PYPI_INDEX_URL} python${PYTHON_VERSION} -m pip install --no-cache-dir /tmp/mooncake-wheel/*.whl && rm -rf /tmp/mooncake-wheel /root/.cache/pip CMD ["/bin/bash"] diff --git a/docs/AGENTS.md b/docs/AGENTS.md new file mode 100644 index 0000000000..debb165ca2 --- /dev/null +++ b/docs/AGENTS.md @@ -0,0 +1,90 @@ +# AGENTS.md - Mooncake Documentation + +This file gives coding agents the repo-local rules for modifying files under +`docs/`. Keep `README.md` as the human-facing quickstart. Use this file for +agent workflow, verification, and maintenance guidance. + +## Scope + +- Applies to changes under `docs/`, especially `docs/source/`. +- Prefer small, reviewable documentation changes. +- Do not rewrite unrelated pages, generated files, or formatting-only content. +- Preserve existing documentation structure unless the user asks for a broader + reorganization. + +## Build and Preview + +Run documentation commands from the `docs` directory: + +``` +cd docs +``` + +Install dependencies with `uv` when needed. The requirements file is +`docs/requirements-docs.txt`; after `cd docs`, use the local filename. If there +is an existing venv, prefer using the existing one. Otherwise, create one before +installing dependencies: + +``` +uv venv +uv pip install -r requirements-docs.txt +``` + +Clean stale build output when validating navigation, generated API pages, or +theme behavior: + +``` +make clean +``` + +Build HTML before handing off user-visible documentation changes: + +``` +make html +``` + +Set `locale` correctly before building when the change depends on localized +content or translated output. + +Serve the generated site for review: + +``` +python -m http.server -d build/html/ +``` + +The default URL is `http://localhost:8000`. If port 8000 is busy, choose another +available port. + +## Editing Guidance + +- Source pages live under `docs/source/`. +- Keep links relative and Sphinx-compatible unless an external URL is required. +- For navigation changes, inspect `docs/source/index.md` and the relevant + toctree before editing individual pages. +- Use Sphinx-native structure for documentation behavior. Do not use client-side + JavaScript or post-render DOM patches for navigation or theme behavior. +- If the requested behavior is not supported by Sphinx or the active theme, + prefer adding a Sphinx extension/plugin instead of patching rendered HTML. +- When a page should be linked from content but excluded from the main sidebar, + use a content link plus appropriate Sphinx metadata such as `orphan: true` + instead of hiding rendered sidebar nodes. +- Keep homepage toctree depth conservative. Do not increase `index.md` maxdepth + unless the user explicitly asks for deeper landing-page nesting. + +## Validation Checklist + +- Run `make html` for docs changes that affect rendered pages, navigation, + cross-references, or Sphinx configuration. +- Check the generated HTML for the changed pages. +- For sidebar or toctree changes, verify both the article body and left sidebar + render the intended entries. +- If a local preview server is useful for review, start one from `docs/` with + `python -m http.server -d build/html/` or an alternate port. + +## Pull Request Hygiene + +- Keep docs-only changes narrowly scoped. +- Review `git diff` before staging so generated files or hook-only formatting + changes do not leak into the PR. +- If opening a PR, use the repository pull request template. +- Use the repository PR title prefix rules from the root `AGENTS.md`. diff --git a/docs/README.md b/docs/README.md index 636b023f98..b3f6798dbe 100644 --- a/docs/README.md +++ b/docs/README.md @@ -11,7 +11,7 @@ cd docs - Install the dependencies: ```bash -pip install -r ../requirements_docs.txt +pip install -r requirements-docs.txt ``` - Clean the previous build (optional but recommended): diff --git a/requirements_docs.txt b/docs/requirements-docs.txt similarity index 91% rename from requirements_docs.txt rename to docs/requirements-docs.txt index 9c267edace..66a4886911 100644 --- a/requirements_docs.txt +++ b/docs/requirements-docs.txt @@ -4,6 +4,8 @@ sphinx-book-theme==1.1.4 sphinx-copybutton==0.5.2 sphinx-design==0.6.1 sphinx-togglebutton==0.3.2 +sphinx-reredirects==0.1.6 +sphinxcontrib-mermaid==2.0.1 myst-parser==3.0.1 # `myst-parser==4.0.1` breaks inline code in titles msgspec snowballstemmer<3 # https://github.com/snowballstem/snowball/issues/229 @@ -16,4 +18,4 @@ git+https://github.com/hmellor/sphinx-autodoc2.git # sphinx-autodoc2==0.5.0 # packages to install to build the documentation cachetools -f https://download.pytorch.org/whl/cpu -torch \ No newline at end of file +torch diff --git a/docs/source/api-reference/cpp/index.md b/docs/source/api-reference/cpp/index.md new file mode 100644 index 0000000000..022c3a08b8 --- /dev/null +++ b/docs/source/api-reference/cpp/index.md @@ -0,0 +1,16 @@ +# C/C++ API Reference + +| Module | Description | +|--------|-------------| +| [Transfer Engine C++ API](transfer-engine) | `TransferEngine` class — memory registration, batch transfer, segment management, RDMA transport | +| [TENT C++ API](tent) | `mooncake::tent::TransferEngine` — next-gen transfer engine with automatic transport selection and fault tolerance | +| [Mooncake Store Client C++ API](mooncake-store) | `Client` class — `Put`/`Get`/`Remove`/`Replicate` operations | + +:::{toctree} +:maxdepth: 1 +:hidden: + +transfer-engine +tent +mooncake-store +::: diff --git a/docs/source/api-reference/cpp/mooncake-store.md b/docs/source/api-reference/cpp/mooncake-store.md new file mode 100644 index 0000000000..49d9ca4ed8 --- /dev/null +++ b/docs/source/api-reference/cpp/mooncake-store.md @@ -0,0 +1,224 @@ +# Mooncake Store C++ API Reference + +## Client C++ API + +### Constructor and Initialization `Init` + +```C++ +ErrorCode Init(const std::string& local_hostname, + const std::string& metadata_connstring, + const std::string& protocol, + void** protocol_args, + const std::string& master_server_entry); +``` + +Initializes the Mooncake Store client. The parameters are as follows: +- `local_hostname`: The `IP:Port` of the local machine or an accessible domain name (default value used if port is not included) +- `metadata_connstring`: The address of the metadata service (e.g., etcd/Redis) required for Transfer Engine initialization +- `protocol`: The protocol supported by the Transfer Engine, including RDMA and TCP +- `protocol_args`: Protocol parameters required by the Transfer Engine +- `master_server_entry`: The address information of the Master (`IP:Port` for default mode and `etcd://IP:Port;IP:Port;...;IP:Port` for high availability mode) + +### Get + +```C++ +tl::expected Get(const std::string& object_key, + std::vector& slices); +``` + +`Get` retrieves the value of `object_key` into the provided `slices`. The returned data is guaranteed to be complete and correct. Each slice must reference local DRAM/VRAM memory that has been pre-registered with `registerLocalMemory(addr, len)` (not the global segments that contribute to the distributed memory pool). The master returns the readable replica list and the client selects a complete replica. Depending on the selected replica, the data may be read from memory, NoF SSD, legacy shared-filesystem `DISK`, client-owned `LOCAL_DISK`, or the configured descriptor-based DFS backend. + +### Put + +```C++ +tl::expected Put(const ObjectKey& key, + std::vector& slices, + const ReplicateConfig& config); +``` + +`Put` stores the value associated with `key` in the configured replica tiers. The `config` parameter controls the number of memory, NoF, and DFS replicas as well as placement preferences. Legacy `DISK` persistence and client-owned `LOCAL_DISK` SSD offload remain asynchronous. When `dfs_replica_num` is `1`, `Put` waits for the requested DFS `WriteAt` operation before returning success; this does not provide an additional `fsync` durability guarantee. + +**Memory Replication Guarantees and Best Effort Behavior:** +- Each slice of an object is guaranteed to be replicated to different segments, ensuring distribution across separate storage nodes +- Different slices from different objects may be placed in the same segment +- Replication operates on a best-effort basis: if insufficient space is available for all requested replicas, the object will still be written with as many replicas as possible + +Requests with `dfs_replica_num == 1` use reliable multi-replica mode: allocation and every requested transfer must succeed, otherwise `Put` fails and allocated replicas are revoked. + +```{warning} +Descriptor-based DFS is a work-in-progress feature for development and evaluation. It is not covered by the Store's production fault-tolerance, HA, durability, or multi-tenant guarantees. +``` + +The DFS-related replica-count fields of `ReplicateConfig` are as follows: + +```C++ +struct ReplicateConfig { + size_t replica_num{1}; // Memory replicas + size_t nof_replica_num{0}; // NoF SSD replicas + size_t dfs_replica_num{0}; // Shared DFS replicas (0 or 1) + SoftPinAction soft_pin_action{SoftPinAction::PRESERVE}; + std::optional soft_pin_ttl_ms{}; // ENABLE override; omitted uses the Master default + bool with_hard_pin{false}; // Whether to enable hard pin (never evicted) + std::string preferred_segment{}; // Preferred segment for allocation + // Other placement, data-type, and grouping fields are omitted. +}; +``` + +`dfs_replica_num` may currently be `0` or `1`. When it is `1`, `replica_num >= 1` is required, so DFS-only placement is not supported. DFS replicas currently support only the `default` tenant and require the master and client DFS backends to be configured with the same shared root and shard layout. See the {ref}`DFS deployment documentation ` for setup and lifecycle limitations. Native C++ clients must initialize a `DistributedStorageBackend` and attach it with `SetDfsStorageBackend()` before issuing DFS reads or writes. DFS descriptors are carried by each `PutStart`, `UpsertStart`, or query response; there is no client-side descriptor cache. The Python/RealClient setup path attaches the backend through `FileStorage`. + +Soft pinning starts when the first replica becomes readable and has a fixed +lifetime: reads do not extend it. `PRESERVE` keeps the committed deadline on an +Upsert, `ENABLE` starts a new lifetime, and `DISABLE` removes it when the write +commits. `soft_pin_ttl_ms` is valid only with `ENABLE`; zero commits ordinary +cache, and values above the Master's configured maximum are rejected. +Soft-pin state is not persisted in snapshots or the HA OpLog; after recovery or +Standby promotion, restored objects are ordinary cache until a later write +explicitly enables soft pinning again. + +### Upsert + +```C++ +tl::expected Upsert(const ObjectKey& key, + std::vector& slices, + const ReplicateConfig& config); + +std::vector> BatchUpsert( + const std::vector& keys, + std::vector>& 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. If either the existing object or the new request has +a DFS replica, a same-size update requires the requested memory, NoF, and DFS +replica counts to match the existing topology. A different-size update releases +the old placement and allocates a new topology. `BatchUpsert` performs the same +operation for multiple keys using a shared replication configuration. + +### Remove + +```C++ +tl::expected Remove(const ObjectKey& key); +``` + +Used to delete the object corresponding to the specified key. This interface marks all data replicas associated with the key in the storage engine as deleted, without needing to communicate with the corresponding storage node (Client). + +### CreateCopyTask + +```C++ +tl::expected CreateCopyTask( + const std::string& key, + const std::vector& targets); +``` + +`CreateCopyTask` creates an asynchronous copy task that will be executed by the client's task execution system. This is useful when you want to submit multiple copy operations without waiting for each one to complete. The task is submitted to the master service, assigned a unique task ID, and executed asynchronously by an available client. The task status can be queried using `QueryTask`. + +**Task Execution and Result Reporting:** +1. **Task Assignment**: The master service assigns the task to an available client during the client's periodic ping operation +2. **Task Execution**: The assigned client executes the copy operation asynchronously in a background thread pool +3. **Result Reporting**: Upon completion (success or failure), the client automatically reports the result to the master service via `MarkTaskToComplete`: + - On success: `status = SUCCESS`, `message = "Task completed successfully"` + - On failure: `status = FAILED`, `message = ` +4. **Status Query**: You can query the task status at any time using `QueryTask` to monitor progress + +**Failure and retry behavior:** +- New tasks enter the queue as `PENDING` and are moved to `PROCESSING` when a client starts execution. +- The client only retries failures caused by `ErrorCode::NO_AVAILABLE_HANDLE`, up to the master-side `max_retry_attempts` limit. +- Submission can fail immediately with `ErrorCode::TASK_PENDING_LIMIT_EXCEEDED` when the task queue is full. + +### CreateMoveTask + +```C++ +tl::expected CreateMoveTask( + const std::string& key, + const std::string& source, + const std::string& target); +``` + +`CreateMoveTask` creates an asynchronous move task that will be executed by the client's task execution system. This is useful when you want to submit multiple move operations without waiting for each one to complete. The task is submitted to the master service, assigned a unique task ID, and executed asynchronously by an available client. The task status can be queried using `QueryTask`. + +**Task Execution and Result Reporting:** +1. **Task Assignment**: The master service assigns the task to an available client during the client's periodic ping operation +2. **Task Execution**: The assigned client executes the move operation asynchronously in a background thread pool +3. **Result Reporting**: Upon completion (success or failure), the client automatically reports the result to the master service via `MarkTaskToComplete`: + - On success: `status = SUCCESS`, `message = "Task completed successfully"` + - On failure: `status = FAILED`, `message = ` +4. **Status Query**: You can query the task status at any time using `QueryTask` to monitor progress + +**Failure and retry behavior:** +- Move tasks follow the same state machine as copy tasks: `PENDING -> PROCESSING -> SUCCESS/FAILED`. +- **Submission-time failures**: If the object or source replica is already missing when `CreateMoveTask` is called, the call returns `ErrorCode::OBJECT_NOT_FOUND` or `ErrorCode::INVALID_PARAMS` directly and no task is created. +- **Execution-time failures**: If the object or source replica disappears after the task is successfully submitted, the task transitions to `FAILED` state. Only `ErrorCode::NO_AVAILABLE_HANDLE` execution failures are retried automatically. +- Timeout and retry policy is configured on the master, not per request. + +### QueryTask + +```C++ +tl::expected QueryTask(const UUID& task_id); +``` + +`QueryTask` queries the status of an asynchronous task (copy or move). This allows you to monitor the progress of task-based operations. The response includes task status, type, creation time, last update time, assigned client, and status message. + +The data structure details of `QueryTaskResponse` are as follows: + +```C++ +struct QueryTaskResponse { + UUID id; // Task UUID + TaskType type; // Task type (REPLICA_COPY or REPLICA_MOVE) + TaskStatus status; // Task status (PENDING, PROCESSING, SUCCESS, or FAILED) + int64_t created_at_ms_epoch; // Task creation timestamp in milliseconds + int64_t last_updated_at_ms_epoch; // Last update timestamp in milliseconds + UUID assigned_client; // UUID of the client assigned to execute the task + std::string message; // Status message or error description +}; +``` + +Typical query-time failure: +- `ErrorCode::TASK_NOT_FOUND`: the task ID is unknown or the completed task has already been evicted from the master's retained finished-task history. + +**Master-side task manager settings affecting task APIs:** +- `--max_total_finished_tasks`: number of completed tasks retained for subsequent `QueryTask` calls +- `--max_total_pending_tasks`: maximum queued tasks before `CreateCopyTask`/`CreateMoveTask` fail with `TASK_PENDING_LIMIT_EXCEEDED` +- `--max_total_processing_tasks`: cap on concurrently processing tasks +- `--pending_task_timeout_sec`: timeout for tasks that remain in `PENDING` (`0` disables it) +- `--processing_task_timeout_sec`: timeout for tasks that remain in `PROCESSING` (`0` disables it) +- `--max_retry_attempts`: retry limit for `NO_AVAILABLE_HANDLE` execution failures + +### BatchQueryIp + +```C++ +tl::expected, boost::hash>, ErrorCode> +BatchQueryIp(const std::vector& client_ids); +``` + +Used to batch query the IP addresses for multiple client IDs. For each client ID in the input list, this interface retrieves the unique IP addresses from all segments mounted by that client. The operation is performed on the Master Service and returns a map from client ID to their IP address lists. Only client IDs that have successfully mounted segments are included in the result map. This is useful for discovering the network locations of storage nodes in the cluster. + +### BatchReplicaClear + +```C++ +tl::expected, ErrorCode> +BatchReplicaClear(const std::vector& object_keys, + const UUID& client_id, + const std::string& segment_name); +``` + +Used to batch clear replicas for multiple object keys belonging to a specific client ID. This interface allows clearing replicas either on a specific segment or across all segments. If segment_name is empty, all replicas of the specified objects are cleared (the objects are deleted entirely). If segment_name is provided, only replicas located on that specific segment are cleared. The operation is performed on the Master Service and returns a list of object keys that were successfully cleared. Only objects that belong to the specified `client_id`, have expired leases, and meet the clearing criteria are processed. This is useful for managing storage resources and cleaning up data on specific storage nodes. + +### QueryByRegex + +```C++ +tl::expected>, ErrorCode> +QueryByRegex(const std::string& str); +``` + +Used to query the replica information for all objects whose keys match the given regular expression. This is useful for batch operations or for retrieving a group of related objects. The operation is performed on the Master and returns a map of keys to their replica lists. + +### RemoveByRegex + +```C++ +tl::expected RemoveByRegex(const ObjectKey& str); +``` + +Used to delete all objects from the store whose keys match the specified regular expression. This provides a powerful way to perform bulk deletions. The command returns the number of objects that were successfully removed. diff --git a/docs/source/design/tent/cpp-api.md b/docs/source/api-reference/cpp/tent.md similarity index 85% rename from docs/source/design/tent/cpp-api.md rename to docs/source/api-reference/cpp/tent.md index 319ae92a92..79ce0278d4 100644 --- a/docs/source/design/tent/cpp-api.md +++ b/docs/source/api-reference/cpp/tent.md @@ -5,13 +5,13 @@ This page summarizes the C++ APIs in `mooncake-transfer-engine/tent/include/tent/transfer_engine.h`. It follows the same structure as the Transfer Engine API documentation. -For conceptual background, see [TENT Overview](overview.md). +For conceptual background, see [TENT Overview](../../design/tent/overview.md). **Prerequisites and API modes** - **Build** with `-DUSE_TENT=ON` to enable TENT. - TENT provides two API surfaces: - **TENT-native API** (documented below). - - **TE-compatible API** via the compatibility shim (set `MC_USE_TENT=1`). For TE-compatible API, see [TE C++ API Reference](../transfer-engine/cpp-api.md) + - **TE-compatible API** via the compatibility shim (set `MC_USE_TENT=1`). For TE-compatible API, see [TE C++ API Reference](transfer-engine.md) **Core APIs vs Advanced APIs** - Core APIs form the minimal path to move data: create the engine, register memory, open segments, submit transfers, and query status. @@ -75,6 +75,7 @@ For users migrating from Transfer Engine, the following table shows how TE APIs | `allocateBatchID(batch_size)` | `allocateBatch(batch_size)` | Renamed | | `freeBatchID(batch_id)` | `freeBatch(batch_id)` | Renamed | | `submitTransfer(batch_id, entries)` | `submitTransfer(batch_id, request_list)` | `TransferRequest` → `Request` | +| *Not available* | `cancelTransfer(batch_id, task_id)` | TENT-only: best-effort cancellation for queued and RDMA tasks | | `submitTransferWithNotify(batch_id, entries, notify_msg)` | `submitTransfer(batch_id, request_list, notifi)` | Unified API with optional notification | | `getTransferStatus(batch_id, task_id, status)` | `getTransferStatus(batch_id, task_id, status)` | Same | | `getBatchTransferStatus(batch_id, status)` | `getTransferStatus(batch_id, status)` | Overloaded; single `TransferStatus` output = overall status | @@ -105,6 +106,64 @@ export MC_USE_TENT=1 When this variable is set, the `mooncake::TransferEngine` class internally delegates to `mooncake::tent::TransferEngine`. Most TE APIs are translated automatically. APIs that have no TENT equivalent (e.g., `installTransport`, `getMetadata`) become no-ops or return placeholder values. +Passing a NIC priority matrix through `installTransport(..., args)` is **not** supported under TENT. Configure custom topology via `MC_TENT_CONF` or `MC_CUSTOM_TOPO_JSON` instead (see below). + +(custom-nic-priority-matrix)= +### Custom NIC Priority Matrix + +TENT accepts the same classic Transfer Engine priority-matrix JSON format: + +```json +{ + "cpu:0": [["mlx5_0"], ["mlx5_1"]], + "cpu:1": [["mlx5_1"], ["mlx5_0"]], + "cuda:0": [["mlx5_0"], ["mlx5_1"]] +} +``` + +Keys must match discovered location names (`cpu:N`, `cuda:N`, `hip:N` on AMD). +TENT accepts the legacy AMD prefix `rocm:N` when parsing matrices so existing +configs keep working — such keys are canonicalized to `hip:N` on load, and +specifying both `rocm:N` and `hip:N` for the same device is rejected as a +conflict. Discovery and dumps emit `hip:N`, matching classic TE. + +Preferred HCAs map to topology rank 0; available/fallback HCAs map to rank 1. + +Configuration priority (highest first): + +1. Inline matrix in `MC_TENT_CONF` / `Config`: `topology/priority_matrix` +2. File path: `topology/custom_json_path` (also set by `MC_CUSTOM_TOPO_JSON`) +3. Automatic topology discovery + +**Inline example (`MC_TENT_CONF`):** + +```json +{ + "topology": { + "priority_matrix": { + "cpu:0": [["mlx5_0"], ["mlx5_1"]], + "cuda:0": [["mlx5_0"], ["mlx5_1"]] + } + } +} +``` + +**Path example:** + +```json +{ + "topology": { + "custom_json_path": "/etc/mooncake/nic_priority_matrix.json" + } +} +``` + +```bash +export MC_CUSTOM_TOPO_JSON=/etc/mooncake/nic_priority_matrix.json +``` + +A topology file may also use TENT's native `{"nics":[...],"mems":[...]}` format when loaded via `custom_json_path` / `MC_CUSTOM_TOPO_JSON`. If loading or parsing fails, TENT falls back to auto-discovery. + ## Core APIs ### Core Usage Path (C++) @@ -189,6 +248,9 @@ struct Request { SegmentID target_id; uint64_t target_offset; size_t length; + int priority = PRIO_HIGH; + std::optional policy_name; + TransportType transport_hint = UNSPEC; }; ``` @@ -197,6 +259,9 @@ struct Request { - `target_id`: Segment ID obtained from `openSegment`. - `target_offset`: Offset within the target segment. - `length`: Number of bytes to transfer. +- `priority`: Scheduling priority. Used by the QoS layer; see [qos.md](../../design/tent/qos.md). +- `policy_name`: Optional. When set, the request matches the named entry instead of the first-matching policy. +- `transport_hint`: Optional. `UNSPEC` (default) defers to `TransportSelector`. Any other `TransportType` pins this request onto that transport for its first try. #### TransferStatus @@ -268,6 +333,24 @@ Queries the status of transfer requests. - `status` / `status_list` / `overall_status`: Output parameter(s) for status. - Return value: `Status::OK()` on success; otherwise a non-OK status. +#### TransferEngine::cancelTransfer + +```cpp +Status cancelTransfer(BatchID batch_id, size_t task_id); +``` + +Requests best-effort cancellation of one public task. A task still waiting in +the TENT admission queue becomes `CANCELED` without being dispatched. For RDMA, +workers suppress slices they observe before `ibv_post_send`; work already +posted to a QP is allowed to drain and may complete successfully. Consequently, +the API returning `OK` means the cancellation request was accepted, not that +the task is already terminal. Continue polling `getTransferStatus` before +calling `freeBatch`. + +Cancellation is idempotent. Merged public tasks share one physical transfer, +so canceling any alias cancels the shared task. Direct cancellation of staging +or non-RDMA transport work currently returns `Status::NotImplemented`. + #### TransferEngine::freeBatch ```cpp @@ -516,13 +599,14 @@ using Location = std::string; const static std::string kWildcardLocation = "*"; ``` -Location strings identify device affinity: `"cpu:0"`, `"cuda:0"`, `"cuda:1"`, etc. Use `"*"` for automatic detection. +Location strings identify device affinity: `"cpu:0"`, `"cuda:0"`, `"hip:0"`, etc. Use `"*"` for automatic detection. On AMD GPUs the canonical prefix is `hip:` (same as classic TE). The legacy TENT prefix `rocm:` is still accepted when parsing locations and custom NIC matrices. ### TransportType ```cpp enum TransportType { - RDMA = 0, + UNSPEC = 0, + RDMA, MNNVL, SHM, NVLINK, @@ -530,7 +614,7 @@ enum TransportType { IOURING, TCP, AscendDirect, - UNSPEC + SUNRISE_LINK, }; ``` diff --git a/docs/source/design/transfer-engine/cpp-api.md b/docs/source/api-reference/cpp/transfer-engine.md similarity index 98% rename from docs/source/design/transfer-engine/cpp-api.md rename to docs/source/api-reference/cpp/transfer-engine.md index eee3629cbb..d1fa4f8838 100644 --- a/docs/source/design/transfer-engine/cpp-api.md +++ b/docs/source/api-reference/cpp/transfer-engine.md @@ -182,6 +182,8 @@ Registers a space starting at address `addr` with a length of `length` on the lo - `length`: The length of the registration space; - `location`: The `device` corresponding to this memory segment, such as `cuda:0` indicating the GPU device, `cpu:0` indicating the CPU socket, by matching with the network card priority order table (see `installTransport`), the preferred network card is identified. You can also use `*`, Transfer Engine will try to automatically recognize the `device` corresponding to `addr`, if it fails to recognize the device, it will print a `WARNING` level log and use all network cards, no preferred network cards. - `remote_accessible`: Indicates whether this memory can be accessed by remote nodes. + For RDMA, `false` registers the buffer for local transfer use only: remote + read/write permissions are not granted and no `rkey` is published. - `update_metadata`: Whether to publish the registration to the metadata service. - Return value: If successful, returns 0; otherwise, returns a negative value. @@ -337,7 +339,7 @@ The HTTP server should implement three following RESTful APIs, while the metadat 2. `PUT /metadata?key=$KEY`: Update the metadata corresponding to `$KEY` to the value of the request body. 3. `DELETE /metadata?key=$KEY`: Delete the metadata corresponding to `$KEY`. -For specific implementation, refer to the demo service implemented in Golang at [mooncake-transfer-engine/example/http-metadata-server](../../../mooncake-transfer-engine/example/http-metadata-server). +For specific implementation, refer to the demo service implemented in Golang at [mooncake-transfer-engine/example/http-metadata-server](gh-dir:mooncake-transfer-engine/example/http-metadata-server). ### Initialization diff --git a/docs/source/api-reference/http/conductor-indexer.md b/docs/source/api-reference/http/conductor-indexer.md new file mode 100644 index 0000000000..0abd67e41c --- /dev/null +++ b/docs/source/api-reference/http/conductor-indexer.md @@ -0,0 +1,493 @@ +# Mooncake Conductor Indexer API + +## Overview + +Mooncake Conductor is a KV cache indexer used by routers and gateways to make +cache-aware scheduling decisions. It consumes KV cache events from inference +engines or storage backends, maintains prefix-hit metadata across cache tiers, +and exposes HTTP APIs for service registration and cache-hit queries. + +This document incorporates the latest API direction from: + +- [RFC #1403: Mooncake KV-Store Indexer API Standardization](https://github.com/kvcache-ai/Mooncake/issues/1403) +- [RFC #1527: KV Events API Standardization](https://github.com/kvcache-ai/Mooncake/issues/1527) + +The Conductor can serve multiple model groups in one process. Each query is +scoped by model identity, block size, LoRA identity, tenant isolation, and the +registered instance that can receive traffic. + +## Concepts + +### Storage tiers + +The indexer tracks KV cache availability across three logical tiers: + +- **G1, Device Pool**: Device-resident KV blocks, such as GPU, NPU, HBM, or + other accelerator memory owned by inference engines. +- **G2, Host Pool**: CPU or host DRAM KV blocks, including Mooncake registered + memory pools. +- **G3, Disk Pool**: SSD, 3FS, DFS, NFS, or other disk-backed KV storage. + +The `medium` field identifies the concrete tier or device type. Common values +are `gpu`, `cpu`, and `disk`. Engines may add other values as new media are +supported. + +### Identity dimensions + +KV cache hits are interpreted under the following dimensions: + +| Dimension | Description | +|---|---| +| `model_name` or `model` | Model identifier. KV blocks from different models are incompatible. | +| `block_size` | Number of tokens per KV block. Different block sizes produce different token-to-block mappings. | +| `additional_salt` | Opaque salt used to separate hash namespaces for quantization, model revision, tenant isolation, or other deployment-specific dimensions. | +|`cache_salt`| Ensure cached data blocks are kept separate for different customers| +| `lora_name` | LoRA adapter name. Empty or `null` means the base model. | +| `tenant_id` | Upstream tenant or customer identity. Used for isolation and to keep query output bounded. | +| `instance_id` | Routable API server or engine instance returned by the Indexer API. Routers use this value as the scheduling target. | +| `backend_id` | KV Events identity for the entity that owns the KV blocks. It may be an inference worker, a Mooncake storage daemon, or another cache backend. | +| `medium` | Cache medium where the blocks are present. | +| `dp_rank` | Data-parallel rank that owns or can serve the blocks. | + +`instance_id` and `backend_id` intentionally have different meanings. +`instance_id` is the router-facing target in the Indexer API. `backend_id` is +the event-facing cache owner in the KV Events API. In deployments where cache +storage is decoupled from inference workers, `backend_id` can identify a cache +daemon while `instance_id` still identifies the engine endpoint that receives +requests. + +## Hashing standard + +The standardized event contract recommends **XXH3-64 with seed `S`**. + +- **Local block hash**: + `XXH3(token_bytes_le, S)`, where tokens are little-endian `u32` values + concatenated for one block. +- **Rolling sequence hash**: + The first block uses `seq_hash[0] = local_block_hash[0]`. Each subsequent + block uses: + + ```text + seq_hash[i] = XXH3(seq_hash[i-1]_le || local_block_hash[i]_le, S) + ``` + + Here `||` means byte concatenation, not a logical OR. + +All hashes used by the standardized KV Events API are rolling sequence hashes. +A `seq_hash` identifies the whole prefix up to that block depth, so equal +prefixes produce equal hashes until the first differing block. + +If an engine does not follow the standardized hashing scheme, it must provide +`token_ids` in `stored` events so the consumer can recompute the indexer's hash +representation. + +## HTTP APIs + +### `POST /register` + +Registers a KV event publisher and starts consuming events from it. + +```json +{ + "endpoint": "tcp://1.1.1.1:5557", + "replay_endpoint": "tcp://1.1.1.1:5558", + "type": "vLLM", + "modelname": "deepseek", + "lora_name": "sql-adapter", + "tenant_id": "default", + "instance_id": "vllm-prefill-node1", + "block_size": 128, + "dp_rank": 0, + "additionalsalt": "w8a8" +} +``` + +| Field | Required | Description | +|---|---|---| +| `endpoint` | Yes | ZMQ KV event publisher endpoint. | +| `replay_endpoint` | No | ZMQ replay endpoint used to recover missed events. | +| `type` | Yes | Publisher type, such as `vLLM`, `SGLang`, or `Mooncake`. | +| `modelname` | Yes | Model name for this publisher. This is the HTTP API wire name for `model_name`. | +| `lora_name` | No | LoRA adapter name. Empty or omitted means base model. | +| `tenant_id` | No | Tenant identity. Defaults to `default`. | +| `instance_id` | Yes | Router-facing engine or API server instance identity. | +| `block_size` | Yes | KV block size in tokens. | +| `dp_rank` | Yes | Data-parallel rank for this publisher. | +| `additionalsalt` | No | HTTP API wire name for `additional_salt`. Defaults to an empty string. | + +Successful response: + +```json +{ + "status": "registered successfully", + "instance_id": "vllm-prefill-node1" +} +``` + +### `POST /unregister` + +Stops consuming events for a registered publisher. + +```json +{ + "type": "vLLM", + "modelname": "deepseek", + "lora_name": "sql-adapter", + "tenant_id": "default", + "instance_id": "vllm-prefill-node1", + "block_size": 128, + "dp_rank": 0 +} +``` + +`tenant_id` defaults to `default`. The current implementation removes the +subscription identified by `(instance_id, tenant_id, dp_rank)`. + +Successful response: + +```json +{ + "status": "unregistered successfully", + "removed_instances": ["vllm-prefill-node1|default|0"] +} +``` + +### `POST /query` + +Query cache hits by token IDs. + +```json +{ + "model": "deepseek", + "lora_name": "sql-adapter", + "token_ids": [101, 15, 100, 55, 89], + "tenant_id": "default", + "instance_id": "vllm-prefill-node1", + "block_size": 64, + "cache_salt": "w8a8" +} +``` + +| Field | Required | Description | +|---|---|---| +| `model` | Yes | Model name. | +| `lora_name` | No | LoRA adapter name. Empty or omitted means base model. | +| `lora_id` | No | Deprecated compatibility field. Do not use together with `lora_name`. | +| `token_ids` | Yes | Prompt token IDs. Only complete blocks are considered. | +| `tenant_id` | No | Tenant identity. Defaults to `default`. | +| `instance_id` | No | If set, query one instance. If omitted, query all instances registered under the tenant. | +| `block_size` | Yes | KV block size in tokens. | +| `cache_salt` | No | Query-side hash namespace salt. Corresponds to the event `additional_salt` concept. | + +Response: + +```json +{ + "default": { + "vllm-prefill-node1": { + "longest_matched": 256, + "GPU": 128, + "DP": { + "0": 128, + "1": 256 + }, + "CPU": 256, + "DISK": 0 + } + } +} +``` + +| Field | Description | +|---|---| +| `longest_matched` | Longest continuous prefix hit in tokens across all tracked media and DP ranks for this instance. | +| `GPU`, `CPU`, `DISK` | Matched prefix tokens available on each medium. Names are examples; future media can be added. | +| `DP` | Matched prefix tokens grouped by data-parallel rank. | + +### `POST /query_by_hash` + +Queries cache hits by precomputed rolling sequence hashes. This API avoids +sending long token lists over the network. + +```json +{ + "model": "deepseek", + "lora_name": "sql-adapter", + "seq_hashes": [1234567890, 9876543210], + "tenant_id": "default", + "instance_id": "vllm-prefill-node1", + "block_size": 64, + "cache_salt": "w8a8" +} +``` + +For compatibility with earlier drafts, clients may call the hash list +`block_hash`, but new clients should use `seq_hashes` to make it explicit that +the values are rolling sequence hashes rather than local block hashes. + +The response shape is the same as `/query`. For this API, +`longest_matched = block_size * matched_hash_count`. + + +## KV Events API + +The KV Events API is the wire contract between cache owners and indexers. +Conductor normalizes engine-specific events into this model. + +### Event envelope + +Every standardized event carries the same envelope: + +```json +{ + "event_id": 42, + "timestamp": 1739145600000, + "event_type": "stored", + "model_name": "llama-3.1-8b", + "block_size": 64, + "additional_salt": null, + "lora_name": null, + "tenant_id": "default", + "backend_id": "worker-0", + "medium": "gpu", + "dp_rank": 0 +} +``` + +| Field | Type | Description | +|---|---|---| +| `event_id` | `u64` | Monotonically increasing sequence number scoped by the event stream dimensions. Authoritative for ordering. | +| `timestamp` | `u64 or null` | Unix epoch milliseconds. Informational only, not used for ordering. | +| `event_type` | `string` | One of `stored`, `removed`, or `cleared`. | +| `model_name` | `string or null` | Model identifier. | +| `block_size` | `u32 or null` | Tokens per block. | +| `additional_salt` | `string or null` | Opaque deployment salt or namespace. | +| `lora_name` | `string or null` | LoRA adapter name, or `null` for the base model. | +| `tenant_id` | `string` | Tenant or customer identity. | +| `backend_id` | `string` | Entity that owns the KV blocks. This can be an engine worker or a decoupled cache daemon. | +| `medium` | `string or null` | Cache medium such as `gpu`, `cpu`, or `disk`. | +| `dp_rank` | `u32 or null` | Data-parallel rank. | + +Events must be processed in consecutive `event_id` order within each stream +identified by `(model_name, block_size, additional_salt, lora_name, tenant_id, +backend_id, medium, dp_rank)`. + +### `stored` + +Published when one or more consecutive blocks are committed to a KV cache. + +```json +{ + "event_id": 42, + "timestamp": 1739145600000, + "event_type": "stored", + "model_name": "llama-3.1-8b", + "block_size": 64, + "additional_salt": null, + "lora_name": null, + "tenant_id": "default", + "backend_id": "worker-0", + "medium": "gpu", + "dp_rank": 0, + "seq_hashes": [1234567890, 9876543210, 1122334455], + "base_block_idx": 5, + "parent_hash": 9999999999, + "token_ids": null +} +``` + +| Field | Type | Description | +|---|---|---| +| `seq_hashes` | `u64[]` | Rolling sequence hashes of consecutive stored blocks. | +| `base_block_idx` | `u32 or null` | Zero-based depth of the first block in this event. | +| `parent_hash` | `u64 or null` | Rolling sequence hash at depth `base_block_idx - 1`; `null` at the root. | +| `token_ids` | `u32[] or null` | Tokens across all blocks in this event. Required when the publisher does not use the standardized hash. | + +At least one of `base_block_idx` or `parent_hash` must be present so the +consumer can locate the blocks in the sequence. + +### `removed` + +Published when one or more blocks are evicted. + +```json +{ + "event_id": 43, + "timestamp": 1739145601000, + "event_type": "removed", + "model_name": "llama-3.1-8b", + "block_size": 64, + "additional_salt": null, + "lora_name": null, + "tenant_id": "default", + "backend_id": "worker-0", + "medium": "gpu", + "dp_rank": 0, + "seq_hashes": [1122334455], + "base_block_idx": 7 +} +``` + +`seq_hashes` is required. `base_block_idx` is optional but recommended for +collision detection and observability. + +### `cleared` + +Published when all blocks for the event stream dimensions are purged. + +```json +{ + "event_id": 44, + "timestamp": 1739145602000, + "event_type": "cleared", + "model_name": "llama-3.1-8b", + "block_size": 64, + "additional_salt": null, + "lora_name": null, + "tenant_id": "default", + "backend_id": "worker-0", + "medium": "gpu", + "dp_rank": 0 +} +``` + +No additional payload fields are required. + +## Compatibility notes + +The current Conductor implementation consumes vLLM ZMQ msgpack batches and +normalizes `BlockStored` and `BlockRemoved` into the internal prefix index. +Registration metadata supplies fields such as `modelname`, `tenant_id`, +`instance_id`, `block_size`, and `additionalsalt` when the engine event does +not carry the full standardized envelope. + +(mooncake-store-master-publisher)= +### Mooncake Store master publisher + +`mooncake_master` can optionally publish RFC #1527 events when +`enable_kv_events=true`. The publisher binds a ZMQ PUB socket +(`kv_events_bind_endpoint`) and emits the same three-frame batch format used by +vLLM/SGLang: empty topic, big-endian sequence number, and a msgpack payload +`[timestamp, [events], dp_rank]`. + +**Per-block events, not global metadata.** Per the +[Dynamo KV Events for Custom Engines](https://docs.nvidia.com/dynamo/kv-managers/kv-events-for-custom-engines) +model, each event describes one or more **KV cache blocks** (`seq_hashes`, +`token_ids`, `parent_hash`, eviction hashes). The master emits **one event per +Mooncake object key** on `PutEnd` / `Remove` / eviction — each key is treated as +one pooled block. Block identity comes from the object key (`seq_hashes` when +the key is decimal/`0x` u64, else `object_key`) and per-object `tenant_id` / +`medium`. The master does **not** stamp process-wide `model_name`, `block_size`, +`lora_name`, or `dp_rank` on events; register those dimensions with the indexer +via `POST /register` (same as decoupled SGLang + storage pool deployments). + +Publisher-level config is limited to transport and stream identity: +`kv_events_bind_endpoint`, `kv_events_backend_id`, and optional compat flags +(`kv_events_emit_object_key`, `kv_events_emit_legacy_compat`). Legacy master +flags such as `kv_events_model_name` are retained for compatibility but are not +written into event payloads. + +Each event map uses RFC #1527 field names (`event_type`, `seq_hashes`, +`backend_id`, `medium`, and so on). When `kv_events_emit_object_key` is enabled +(default), the map also includes `object_key` with the Mooncake store key so +Dynamo and other consumers can match on `sha256` + Mooncake key format without +requiring decimal/`0x` `seq_hash` encoding. When `kv_events_emit_legacy_compat` +is enabled (default), the map also includes vLLM-compatible aliases such as +`type` and `block_hashes` so Dynamo relay mode can forward events without an +adapter. + +Object keys may encode the rolling `seq_hash` as a decimal or `0x`-prefixed +hex string; when `seq_hash` cannot be parsed, events are still published if +`kv_events_emit_object_key=true` (with an empty `seq_hashes` array). Configure +`backend_id` to identify the cache owner (for example a per-node storage +daemon) and register the bind endpoint with the indexer using publisher type +`Mooncake`. + +### Field provenance matrix (SGLang vs master vs indexer registration) + +In decoupled deployments (inference workers + Mooncake host/disk pool), the +global KV indexer merges **three sources of truth**. Use this table when +splitting publishers or writing PR/integration notes. + +**Legend** + +| Symbol | Meaning | +|---|---| +| **SGLang** | Inference engine ZMQ KV events (`BlockStored` / `BlockRemoved` / `AllBlocksCleared`) | +| **Master** | `mooncake_master` optional RFC #1527 publisher (`enable_kv_events`) | +| **Register** | Indexer HTTP `POST /register` (or CLI `--workers`) — not carried on the event wire | +| **S+M** | Either source may supply; must agree on value for the stream | +| **—** | Not applicable for that event type | + +#### Envelope and stream identity + +| Field | SGLang | Master | Register | Notes | +|---|---|---|---|---| +| `event_id` | Yes | Yes | — | Each publisher maintains its own monotonic counter per stream. | +| `timestamp` | Yes | Yes | — | Informational only; not used for ordering. | +| `event_type` | Yes | Yes | — | `stored` / `removed` / `cleared`. | +| `model_name` | S+M | — | S+M | Register uses `modelname`. Engine events carry per-block context; master omits (nil). | +| `block_size` | Yes | — | Yes | Required for token↔block mapping. Register supplies for master publisher. | +| `additional_salt` | Yes | — | S+M | Register uses `additionalsalt`. Engine per-block; master omits (nil). | +| `lora_name` | Yes | — | S+M | Per-block on engine events; master has no adapter context. | +| `tenant_id` | S+M | Yes | Yes | Per-object on master events. Register default `default`. | +| `backend_id` | S+M | Yes | — | **Master**: storage daemon / pool owner. **SGLang**: often worker id; in decoupled mode prefer master=`daemon`, engine via **Register** `instance_id`. | +| `medium` | Yes | Partial | — | **SGLang**: `GPU`, `CPU_PINNED`, `DISK`, `EXTERNAL`, etc. **Master**: only `cpu` / `disk` (host/disk pool), never GPU. | +| `dp_rank` | Yes | — | Yes | Per-batch on engine ZMQ wire. Master batch trailer uses `0`; register dp_rank with indexer. | + +#### `stored` payload + +| Field | SGLang | Master | Register | Notes | +|---|---|---|---|---| +| `seq_hashes` | Yes | Conditional | — | **Required from SGLang** for correct prefix index. Master: single hash when key is decimal/`0x` u64; empty array when only `object_key` is used. | +| `object_key` | — | Yes | — | Mooncake store key (`kv_events_emit_object_key`, default on). Used by Dynamo for sha256+key matching. | +| `block_hashes` (legacy) | Yes | Conditional | — | Alias of `seq_hashes` when `kv_events_emit_legacy_compat` is enabled on master. | +| `parent_hash` | Yes | — | — | Radix parent link; master has no sequence tree. | +| `parent_block_hash` (legacy) | Yes | — | — | Same as `parent_hash`. | +| `base_block_idx` | Yes | Partial | — | Depth of first block in batch; master uses `0` for standalone pool blocks. | +| `token_ids` | Yes | — | — | Required for `/query` by tokens or hash recomputation when engine is non-standard. | +| `block_size` (in-event) | Yes | — | — | Per-block token count in SGLang `BlockStored`; master uses envelope-level config only. | + +#### `removed` payload + +| Field | SGLang | Master | Register | Notes | +|---|---|---|---|---| +| `seq_hashes` | Yes | Conditional | — | **Required** on wire for strict RFC consumers. Master emits one hash when parseable, else empty with `object_key`. | +| `base_block_idx` | Yes | — | — | Optional but recommended for observability. | + +#### `cleared` payload + +| Field | SGLang | Master | Register | Notes | +|---|---|---|---|---| +| (no extra fields) | — | — | — | Event is envelope-only. | +| `cleared` / `AllBlocksCleared` | Yes | — | — | Engine `reset()` / full cache flush. Master does not emit today. | + +#### Indexer / router plane (not in KV event JSON) + +| Field | SGLang | Master | Register | Notes | +|---|---|---|---|---| +| `instance_id` | — | — | Yes | Router-facing schedule target. Distinct from `backend_id`. | +| `endpoint` | — | — | Yes | ZMQ PUB to subscribe (SGLang or master bind address). | +| `replay_endpoint` | — | — | Yes | Optional gap replay (engine ROUTER). | +| `type` | — | — | Yes | Publisher kind: `vLLM`, `SGLang`, `Mooncake`, etc. | + +#### Recommended split for Dynamo global KV indexer + +```mermaid +flowchart LR + SGLang["SGLang ZMQ"] + Master["Mooncake master ZMQ"] + Reg["POST /register"] + Idx["Global KV indexer"] + + SGLang -->|"GPU + HiCache tiers
tokens, parent_hash, seq_hashes, lora"| Idx + Master -->|"Host/Disk pool
backend_id, medium=cpu|disk"| Idx + Reg -->|"instance_id, model, block_size"| Idx +``` + +| Capability | Primary source | +|---|---| +| GPU prefix hits, LoRA-aware hashes, parent chain, multi-block batches | **SGLang** | +| Pooled host/disk replica visibility | **Master** (if keys encode `seq_hash`) | +| Request routing target | **Register** (`instance_id`) | +| Tiered `/query` response (`gpu` / `cpu` / `disk`) | Merge **SGLang** + **Master** events (see RFC #1403) | diff --git a/docs/source/api-reference/http/http-service.md b/docs/source/api-reference/http/http-service.md new file mode 100644 index 0000000000..0050adcd50 --- /dev/null +++ b/docs/source/api-reference/http/http-service.md @@ -0,0 +1,518 @@ +# Mooncake Store HTTP Service + +The Mooncake Store HTTP Service provides RESTful endpoints for cluster management, monitoring, and data operations. This service is embedded within the `mooncake_master` process and can be enabled alongside the primary RPC services. + +## Overview + +The HTTP service serves multiple purposes: +- **Metrics & Monitoring**: Prometheus-compatible metrics endpoints +- **Cluster Management**: Query and manage distributed storage segments +- **Data Inspection**: Examine stored objects and their replicas +- **Health Checks**: Service availability and status verification + +The Python `mooncake.mooncake_store_service` module also provides a lightweight +Store REST API for data operations and standalone segment mount/unmount +workflows. Unless configured otherwise, it listens on port `8080`. + +## HTTP Endpoints + +### Metrics Endpoints + +#### `/metrics` +Prometheus-compatible metrics endpoint providing detailed system metrics in text format. + +**Method**: `GET` +**Content-Type**: `text/plain; version=0.0.4` +**Response**: Comprehensive metrics including request counts, error rates, latency statistics, and resource utilization + +**Example**: +```bash +curl http://localhost:8080/metrics +``` + +#### `/metrics/summary` +Human-readable metrics summary with key performance indicators. + +**Method**: `GET` +**Content-Type**: `text/plain; version=0.0.4` +**Response**: Condensed overview of system health and performance metrics + +**Example**: +```bash +curl http://localhost:8080/metrics/summary +``` + +### Data Management Endpoints + +#### `/query_key` +Retrieve replica information for a specific key, including memory locations and transport endpoints. + +**Method**: `GET` +**Parameters**: `key` (query parameter) - The object key to query +**Content-Type**: `application/json; charset=utf-8` +**Response**: JSON object with success status and replica data array + +**Example**: +```bash +curl "http://localhost:8080/query_key?key=my_object" +``` + +**Success Response** (HTTP 200): +```json +{ + "success": true, + "data": [ + { + "size_": 1073741824, + "buffer_address_": 140732000000000, + "protocol_": "rdma", + "transport_endpoint_": "192.168.1.100:12345" + } + ] +} +``` + +**Error Response** (key not found, HTTP 404): +```json +{ + "success": false, + "error_code": -704, + "error_message": "OBJECT_NOT_FOUND" +} +``` + +**Error Response** (service unavailable, HTTP 503): +```json +{ + "success": false, + "error_code": -1011, + "error_message": "service plane is not active" +} +``` + +#### `/batch_query_keys` +Retrieve replica information for multiple keys in a single request, including memory locations and transport endpoints for each key. The endpoint performs a read-only metadata lookup and does not grant leases, trigger promotion, or update cache-hit metrics. + +**Method**: `GET` +**Parameters**: `keys` (query parameter) - Comma-separated list of object keys to query (format: key1,key2,key3) +**Content-Type**: `application/json; charset=utf-8` +**Response**: JSON-formatted mapping of keys to their respective replica descriptors + +**Example**: +```bash +curl "http://localhost:8080/batch_query_keys?keys=key1,key2,key3" +``` + +**Response Format**: +```text +{ + "success": true, + "data": { + "key1": { + "ok": true, + "values": [ + { + "transport_endpoint_": "hostname:port", + "buffer_descriptor": {...} + } + ], + "disk_values": [ + { + "file_path": "/path/to/object", + "object_size": 4096 + } + ], + "local_disk_values": [ + { + "client_id": "12345-67890", + "object_size": 4096, + "transport_endpoint": "hostname:port" + } + ], + "nof_values": [ + { + "transport_endpoint_": "hostname:port", + "buffer_descriptor": {...} + } + ] + }, + "key2": { + "ok": false, + "error": "error message" + } + } +} +``` + +The `values` field is always present (empty array when no memory replica exists). The `disk_values`, `local_disk_values`, and `nof_values` fields are optional and only appear when the corresponding replica type is present for the key. + +#### `/get_all_keys` +List all keys currently stored in the distributed system. + +**Method**: `GET` +**Content-Type**: `text/plain; version=0.0.4` +**Response**: Newline-separated list of all stored keys + +**Example**: +```bash +curl http://localhost:8080/get_all_keys +``` + +### Segment Management Endpoints + +#### `/get_all_segments` +List all mounted segments in the cluster. + +**Method**: `GET` +**Content-Type**: `text/plain; version=0.0.4` +**Response**: Newline-separated list of segment names + +**Example**: +```bash +curl http://localhost:8080/get_all_segments +``` + +#### `/query_segment` +Query detailed information about a specific segment, including used and available capacity. + +**Method**: `GET` +**Parameters**: `segment` (query parameter) - Segment name to query +**Content-Type**: `text/plain; version=0.0.4` +**Response**: Multi-line text with segment details + +**Example**: +```bash +curl "http://localhost:8080/query_segment?segment=segment_name" +``` + +**Response Format**: +``` +segment_name +Used(bytes): 1073741824 +Capacity(bytes): 4294967296 +``` + +#### `/get_segments_detail` +Get detailed information of all segments in JSON format, including segment metadata, allocator usage, and status. + +**Method**: `GET` +**Content-Type**: `application/json; charset=utf-8` +**Response**: JSON object containing an array of segment details + +**Example**: +```bash +curl http://localhost:8080/get_segments_detail +``` + +**Response Format**: +```json +{ + "total_segments": 2, + "segments": [ + { + "segment_name": "segment_0", + "segment_id": "00000000-0000-0000-0000-000000000001", + "client_id": "00000000-0000-0000-0000-000000000002", + "base_address": "0x300000000", + "size_bytes": 17179869184, + "size_human": "16 GiB", + "te_endpoint": "192.168.1.1:12345", + "protocol": "rdma", + "status": "MOUNTED", + "allocator_used_bytes": 1073741824, + "allocator_capacity_bytes": 17179869184, + "allocator_usage_percent": 6.25 + } + ] +} +``` + +**Fields**: +- `total_segments` (integer): Total number of segments in the cluster +- `segments` (array): Array of segment detail objects + - `segment_name` (string): Name of the segment + - `segment_id` (string): UUID of the segment + - `client_id` (string): UUID of the client that owns the segment + - `base_address` (string): Base memory address in hex + - `size_bytes` (integer): Segment size in bytes + - `size_human` (string): Human-readable segment size + - `te_endpoint` (string): Transport endpoint address + - `protocol` (string): Transfer protocol (e.g., rdma, tcp) + - `status` (string): Current segment status + - `allocator_used_bytes` (integer): Bytes currently allocated + - `allocator_capacity_bytes` (integer): Total allocator capacity in bytes + - `allocator_usage_percent` (number): Percentage of allocator capacity used + +### Health Check Endpoints + +#### `/health` +Basic health check endpoint for service availability verification. + +**Method**: `GET` +**Content-Type**: `text/plain; version=0.0.4` +**Response**: `OK` when service is healthy +**Status Codes**: +- `200 OK`: Service is healthy +- Other: Service may be experiencing issues + +**Example**: +```bash +curl http://localhost:8080/health +``` + +#### `/version` +Report the master version. Always available, including while the master is in +standby. + +**Method**: `GET` +**Content-Type**: `application/json; charset=utf-8` +**Response**: JSON object with: +- `version` (string): Store version used for RPC handshake compatibility +- `display_version` (string): Human-readable release plus short git hash + +**Example**: +```bash +curl http://localhost:8080/version +``` + +```json +{"version":"2.0.0","display_version":"0.3.12.post1 (git: f9e8311f)"} +``` + +Real clients expose the same `/version` payload on their own client HTTP port +when `enable_client_http_server` is on. See +[Client Metrics Endpoint](../../getting_started/observability.md#client-metrics-endpoint). + +## Store REST API Endpoints + +The following endpoints are served by the Python store REST service, which wraps +`MooncakeDistributedStore` with an aiohttp service. The HTTP handlers live in +Python, while mount and unmount operations are delegated to the underlying store +binding. Start the service with: + +```bash +python -m mooncake.mooncake_store_service \ + --config /path/to/mooncake_config.json \ + --port 8080 +``` + +If the wheel console scripts are installed, the equivalent command is: + +```bash +mc_store_rest_server --config /path/to/mooncake_config.json --port 8080 +``` + +### `/api/mount_shm` +Mount a named shared memory object as one or more Mooncake store segments. +Protocols with a registration-size limit split oversized regions and return +multiple segment ids. Protocols without such a Store-level limit, such as TCP +and RDMA, use a single segment regardless of `max_mr_size`. + +**Method**: `POST` +**Content-Type**: `application/json` + +**Request Body**: +```json +{ + "name": "mooncake_segment", + "size": 16777216, + "offset": 0, + "protocol": "tcp", + "location": "" +} +``` + +**Fields**: +- `name` (string, required): Named shared memory object name. A leading `/` is + accepted, but path separators are not. +- `size` (integer, required): Number of bytes to mount. +- `offset` (integer, optional): File offset in bytes. Defaults to `0`. +- `protocol` (string, optional): Transfer protocol. Defaults to the service + configuration protocol. +- `location` (string, optional): Device or locality hint. Defaults to an empty + string. + +**Success Response**: +```json +{ + "status": "success", + "segment_ids": ["00000000-0000-0000-0000-000000000001"] +} +``` + +**Example**: +```bash +curl -X POST http://localhost:8080/api/mount_shm \ + -H "Content-Type: application/json" \ + -d '{ + "name": "mooncake_segment", + "size": 16777216, + "offset": 0, + "protocol": "tcp", + "location": "" + }' +``` + +### `/api/unmount_shm` +Unmount one or more segment ids previously returned by `/api/mount_shm`. + +**Method**: `POST` +**Content-Type**: `application/json` + +**Request Body**: +```json +{ + "segment_ids": ["00000000-0000-0000-0000-000000000001"], + "grace_period_seconds": 0 +} +``` + +`segment_ids` may also be provided as a single string for one segment. +`grace_period_seconds` is optional and defaults to `0`, which keeps the +existing immediate unmount behavior. When set to a positive value, the master +keeps the segment readable for that grace period while preventing new +allocations, then completes the unmount. + +**Success Response**: +```json +{ + "status": "success" +} +``` + +**Example**: +```bash +curl -X POST http://localhost:8080/api/unmount_shm \ + -H "Content-Type: application/json" \ + -d '{"segment_ids": ["00000000-0000-0000-0000-000000000001"], + "grace_period_seconds": 30}' +``` + +### `/api/mount` +Allocate memory inside the store process and mount it as one or more Mooncake +store segments. Protocols with a registration-size limit split oversized +requests and return multiple segment ids. Protocols without such a Store-level +limit, such as TCP and RDMA, use a single segment regardless of `max_mr_size`. +The response includes the actual allocated size after alignment. + +**Method**: `POST` +**Content-Type**: `application/json` + +**Request Body**: +```json +{ + "size": 16777216, + "protocol": "tcp", + "location": "" +} +``` + +**Fields**: +- `size` (integer, required): Number of bytes requested. Must be positive. +- `protocol` (string, optional): Transfer protocol. Defaults to the service + configuration protocol. +- `location` (string, optional): Device or locality hint. Defaults to an empty + string. + +**Success Response**: +```json +{ + "status": "success", + "segment_ids": ["00000000-0000-0000-0000-000000000002"], + "allocated_size": 16777216 +} +``` + +**Example**: +```bash +curl -X POST http://localhost:8080/api/mount \ + -H "Content-Type: application/json" \ + -d '{"size": 16777216, "protocol": "tcp", "location": ""}' +``` + +### `/api/unmount` +Unmount one or more segment ids previously returned by `/api/mount` and free +the memory allocated by the store process. + +**Method**: `POST` +**Content-Type**: `application/json` + +**Request Body**: +```json +{ + "segment_ids": ["00000000-0000-0000-0000-000000000002"], + "grace_period_seconds": 0 +} +``` + +`segment_ids` may also be provided as a single string for one segment. +`grace_period_seconds` is optional and defaults to `0`, which keeps the +existing immediate unmount-and-free behavior. When set to a positive value, the +master keeps the segment readable for that grace period while preventing new +allocations, then the store releases the local allocated memory after cleanup. + +**Success Response**: +```json +{ + "status": "success" +} +``` + +**Example**: +```bash +curl -X POST http://localhost:8080/api/unmount \ + -H "Content-Type: application/json" \ + -d '{"segment_ids": ["00000000-0000-0000-0000-000000000002"], + "grace_period_seconds": 30}' +``` + +### `/api/unmount_local_disk` +Deregister this store's SSD offload tier from the master before the process +goes away. Intended for a shutdown hook. + +The master stops naming this store as the owner of the keys it offloaded, so a +reader gets a clean miss instead of a peer that is about to disappear. Without +this, a `LOCAL_DISK` segment leaves the master only when the client expires — +one `client_ttl` after the store stops pinging — and reads that pick up the +stale owner in that window block on the connect retries (see +`MC_RPC_CONNECT_TIMEOUT_MS`) before missing. + +The call then holds for `grace_period_seconds` before returning. Unlike a memory +replica, which the NIC serves without help from the store process, a disk +replica is read and pushed by that process, so it has to stay alive for the +reads the master handed out before the deregistration. Offloading is stopped for +good when this is called; the store is expected to exit afterwards. + +Returns success and does nothing when SSD offload is not enabled on this store. +Safe to call more than once. + +**Method**: `POST` +**Content-Type**: `application/json` + +**Request Body**: +```json +{ + "grace_period_seconds": 30 +} +``` + +`grace_period_seconds` is optional and defaults to `0`, which returns as soon as +the master has dropped the segment. Must be a non-negative integer no greater +than 3600 (1 hour); a malformed body or an out-of-range value gets a `400` +without touching the store, so a mistake here (seconds where milliseconds were +meant, say) cannot block a preStop hook for hours. + +**Success Response**: +```json +{ + "status": "success" +} +``` + +**Example** — as a Kubernetes preStop hook, with a +`terminationGracePeriodSeconds` longer than the grace period: +```bash +curl -X POST http://localhost:8080/api/unmount_local_disk \ + -H "Content-Type: application/json" \ + -d '{"grace_period_seconds": 30}' +``` diff --git a/docs/source/api-reference/http/index.md b/docs/source/api-reference/http/index.md new file mode 100644 index 0000000000..c67a9fbede --- /dev/null +++ b/docs/source/api-reference/http/index.md @@ -0,0 +1,14 @@ +# HTTP + +| Module | Description | +|--------|-------------| +| [HTTP Service](http-service) | RESTful endpoints for cluster management, metrics, and data inspection | +| [Conductor Indexer API](conductor-indexer) | Cache-aware routing: service registration, KV event subscription, prefix cache-hit query APIs | + +:::{toctree} +:maxdepth: 1 +:hidden: + +http-service +conductor-indexer +::: diff --git a/docs/source/api-reference/python/dataproto-structured-object-transfer.md b/docs/source/api-reference/python/dataproto-structured-object-transfer.md new file mode 100644 index 0000000000..4545d6e1c3 --- /dev/null +++ b/docs/source/api-reference/python/dataproto-structured-object-transfer.md @@ -0,0 +1,172 @@ +# DataProto structured object usage + +Mooncake can store DataProto-like objects as structured objects so callers can pass a lightweight handle between stages and materialize only the fields they need. + +A DataProto-like object is any object with these mapping-like attributes: + +- `batch`: tensor or ndarray fields indexed by batch row. +- `non_tensor_batch`: per-row non-tensor fields. +- `meta_info`: small metadata for the whole batch. + +Plain dictionaries are also accepted. A dictionary with only `batch`, `non_tensor_batch`, and `meta_info` keys is treated as an envelope; other dictionaries are treated as `batch` fields. + +## Public API + +```python +from mooncake.structured_object_store import ( + BundleTransferPolicy, + MooncakeBundleTransfer, + export_dataproto_ref, + import_dataproto_ref, + tensor_object_buffer, +) + +transfer = MooncakeBundleTransfer(store, key_prefix="rl") + +ref = transfer.put_dataproto( + data, + namespace="rollout", + partition="step-1", + stage="rollout", +) + +ref = transfer.append_dataproto_fields( + ref, + logprob_data, + stage="old_log_prob", +) + +subset = transfer.get_dataproto( + ref, + fields=["input_ids", "old_log_probs"], + meta_info_keys=["step"], +) + +handle = export_dataproto_ref(ref) +ref = import_dataproto_ref(handle) +transfer.cleanup_dataproto(ref) +``` + +## Lightweight handles + +`MooncakeDataProtoRef` contains only DataProto-level routing information: + +- `stage_refs`: stage name to structured object reference. +- `field_index`: field name to `(stage, member, section)`. +- `batch_size`, `namespace`, `partition`, `meta_info`, and optional `global_indexes`. + +It does not duplicate dtype, shape, chunk layout, or range metadata. Those details remain in the structured object manifest. Use `dataproto_manifest_view(ref)` when a caller needs an introspection view derived from the manifests. + +For process boundaries, use `export_dataproto_ref(ref)`. The exported handle is JSON-safe and contains manifest keys instead of embedded manifest payloads. `get_dataproto()`, `append_dataproto_fields()`, `dataproto_manifest_view()`, and `cleanup_dataproto()` accept either an in-memory ref or an exported handle. + +## Writing fields + +`put_dataproto()` writes one structured object for the requested stage. `append_dataproto_fields()` writes another structured object and updates the handle. Existing fields are not rewritten. + +Duplicate field names are rejected by default. Use `overwrite=True` only when replacing all fields from an existing stage; after the new stage object is written successfully, the old stage object is removed. + +Field names are global within a ref. A `batch` field and a `non_tensor_batch` field cannot use the same name. + +## Reading fields + +`get_dataproto()` supports: + +- `fields`: mixed batch and non-tensor field selection. +- `batch_fields`: batch-only selection. +- `non_tensor_fields`: non-tensor-only selection. +- `meta_info_keys`: metadata selection. +- `data_cls`: return a DataProto-like class instead of a plain dict. +- `destinations`: caller-provided output buffers. +- `rows`: row selection with a Python `slice`, `StructuredMemberSlice`, or an integer row-index sequence. + +Use `rows` to materialize the same batch rows across all selected `batch` and `non_tensor_batch` fields: + +```python +subset = transfer.get_dataproto( + ref, + fields=["input_ids", "text", "rewards"], + rows=slice(128, 256), +) + +gathered = transfer.get_dataproto( + ref, + batch_fields=["input_ids"], + rows=[7, 3, 9], +) +``` + +Row selection supports `axis=0`. Tensor and ndarray batch fields, including native Mooncake tensor fields, are read by byte ranges from the stored payload. Structured object `non_tensor_batch` fields read only the selected row metadata and payload ranges. `destinations` may be combined with `rows` for fields that support caller-provided output buffers. Use `raw_destination(ptr, size, owner, pre_registered=True)` for BufferPool or otherwise pre-registered destination memory. + +The result is a plain dictionary when `data_cls` is omitted: + +```python +{ + "batch": {...}, + "non_tensor_batch": {...}, + "meta_info": {...}, +} +``` + +If `data_cls` is provided, Mooncake first tries `data_cls.from_dict(batch, non_tensor_batch, meta_info=meta_info)`, then falls back to `data_cls(batch=..., non_tensor_batch=..., meta_info=...)`. + +## Tensor and ndarray behavior + +Tensor fields are stored through the best available Mooncake path: + +1. `tensor_object_buffer(ptr, size, owner, batch_size=...)` with `copy_mode="zero_copy"` uses `put_tensor_from()` directly. +2. Torch tensors use the store tensor API when available. +3. If a tensor-native path is unavailable, Mooncake falls back to a serialized tensor payload. +4. Scalar tensors use a correctness fallback until the native tensor codec preserves zero-dimensional shape. + +Numeric numpy arrays are stored as structured ndarray members. Non-tensor ndarray PUTs are staged and copied before being written to Mooncake; non-contiguous arrays are made contiguous as part of that staging path. Row slices are materialized through structured range reads when the backend supports them, and `destinations` can still be used on GET to materialize selected fields into caller-provided buffers. + +`non_tensor_batch` object arrays are structured-encoded according to their contents. Numeric scalar object arrays, ragged tensors, bytes, strings, JSON-like values, and selected media payloads have explicit codecs. These fields are serialized by design and should not be treated as zero-copy tensor data. + +For typed-ragged ndarray rows, Mooncake copies the rows directly into BufferPool staging memory with the native fast-copy extension and writes each chunk through `batch_put_from`. If the extension or BufferPool path is unavailable, the implementation transparently falls back to the regular copy path. This optimization does not change the stored format or require caller configuration. + +## Materializing into caller buffers + +`destinations` can reuse caller-provided buffers: + +```python +dst = np.empty((rows, width), dtype=np.int64) +result = transfer.get_dataproto(ref, batch_fields=["input_ids"], destinations={"input_ids": dst}) +assert result["batch"]["input_ids"] is dst +``` + +For tensor payloads stored as Mooncake tensor objects or tensor-object buffers, pass a `tensor_object_buffer` destination: + +```python +lease = pool.acquire(nbytes) +result = transfer.get_dataproto( + ref, + batch_fields=["hidden_states"], + destinations={ + "hidden_states": tensor_object_buffer( + lease.ptr, + lease.size, + lease, + batch_size=batch_size, + ) + }, +) +``` + +The destination owner or lease must remain alive for as long as the materialized data may be used. + +## Copy policy + +Control-plane metadata and manifests are small and always use the copy path. Tensor and ndarray payloads use the configured payload policy: + +```python +policy = BundleTransferPolicy(copy_mode="zero_copy") +transfer.put_structured_object(payload, policy=policy) +``` + +`copy_mode="zero_copy"` requires tensor payloads to be provided as `tensor_object_buffer`; plain torch tensors are rejected because they do not expose a registered tensor-object buffer. + +Typed-ragged and other non-tensor structured payloads do not support `copy_mode="zero_copy"`. Use `auto` to enable BufferPool staging and the native fast-copy path when available, or `copy` to force the regular `store.put` path. + +## Cleanup + +`cleanup_dataproto(ref)` removes all stage objects referenced by the handle. It accepts both in-memory refs and exported transport handles. diff --git a/docs/source/api-reference/python/ep-backend.md b/docs/source/api-reference/python/ep-backend.md new file mode 100644 index 0000000000..a7498f10b7 --- /dev/null +++ b/docs/source/api-reference/python/ep-backend.md @@ -0,0 +1,439 @@ +# Mooncake EP & Mooncake PG + +## Overview + +Mooncake provides two closely related components for fault-tolerant MoE +inference: + +- **Mooncake PG** is a `torch.distributed` ProcessGroup backend. It registers + the `mooncake` accelerator backend and the `mooncake-cpu` backend, implements + collective and point-to-point APIs, and exposes dynamic-membership helpers. +- **Mooncake EP** is an expert-parallel dispatch/combine runtime for + latency-sensitive MoE inference. It follows the DeepEP low-latency programming + model while adding rank activeness awareness and Mooncake transport support. + +The usual integration pattern is to initialize a Mooncake process group first, +then construct a Mooncake EP `Buffer` from that group. The process group is used +both for regular collectives and for exchanging EP bootstrap metadata. + +For implementation details, see the +[Mooncake PG design guide](../../design/mooncake-backend-pg.md) and the +[Mooncake EP design guide](../../design/mooncake-ep.md). + +## Installation and build notes + +Mooncake EP and PG are included in CUDA-enabled Mooncake wheels. When building +from source, enable the EP/PG extensions with: + +```bash +cmake .. -DWITH_EP=ON +``` + +The extensions are compiled against a specific PyTorch version. At import time, +`mooncake.pg` and `mooncake.ep` load version-suffixed extension modules that +match the active `torch.__version__`. If the current PyTorch version does not +match a built extension, import will fail with a message such as +`Mooncake PG was not built against torch==...`. + +## Mooncake PG quick start + +### CUDA backend + +```python +import os + +import torch +import torch.distributed as dist +from mooncake import pg + + +rank = int(os.environ["RANK"]) +world_size = int(os.environ["WORLD_SIZE"]) +local_rank = int(os.environ.get("LOCAL_RANK", rank)) + +torch.cuda.set_device(local_rank) +device = torch.device("cuda", local_rank) + +dist.init_process_group( + backend="mooncake", + rank=rank, + world_size=world_size, +) + +x = torch.tensor([rank + 1], dtype=torch.int32, device=device) +dist.all_reduce(x, op=dist.ReduceOp.SUM) +print(f"rank={rank}, all_reduce={int(x.cpu())}") +``` + +Run it with the usual PyTorch launcher, for example: + +```bash +torchrun --nproc-per-node=2 pg_quickstart.py +``` + +### CPU backend + +Use `backend="mooncake-cpu"`: + +```python +dist.init_process_group( + backend="mooncake-cpu", + rank=rank, + world_size=world_size, +) +``` + +`pg_options` is optional for a fixed-size group using the default failure +handling. Pass `MooncakeBackendOptions` when reserving additional group +capacity, joining as an extension, or selecting a non-default failure mode. + +### Selecting network devices + +To explicitly restrict Mooncake to a list of NIC / HCA devices, call +`pg.set_device_filter(...)` before `init_process_group()`: + +```python +from mooncake import pg + +pg.set_device_filter(["mlx5_1", "mlx5_2"]) +``` + +For test and benchmark commands, the same setting is commonly passed through +`MOONCAKE_PGTEST_DEVICE_FILTERS=mlx5_1,mlx5_2`. + +## Mooncake PG Torch API reference + +### `MooncakeBackendOptions` + +```python +pg.MooncakeBackendOptions(max_group_size) +pg.MooncakeBackendOptions(max_group_size, is_extension) +pg.MooncakeBackendOptions( + max_group_size, + is_extension, + auto_deactivate_on_failure, + auto_sync_on_failure, +) + +# Explicit active-rank mirror overloads +pg.MooncakeBackendOptions(active_ranks) +pg.MooncakeBackendOptions(active_ranks, is_extension) +pg.MooncakeBackendOptions(active_ranks, is_extension, max_group_size) +``` + +Arguments: + +- `max_group_size`: fixed in-group slot capacity. It must be at least the + initially declared group size and cannot be increased later. +- `active_ranks`: optional contiguous `torch.int32` storage used as a mirror of + committed PG membership. Its initial contents are ignored. Size it to + `max_group_size`; it may be on CPU or GPU. +- `is_extension`: set to `True` for a replacement or joining process that will + enter an existing group through `join_group()`. +- `auto_deactivate_on_failure` and `auto_sync_on_failure`: select automatic or + framework-managed failure handling. Both default to `True`; auto-sync + requires auto-deactivation. + +### Utility functions + +| Function | Purpose | Notes | +| --- | --- | --- | +| `pg.set_host_ip(host_ip)` | Override the host IP used by the backend. | Call before `init_process_group()`. | +| `pg.set_device_filter(filters)` | Restrict NIC/HCA selection. | Call before `init_process_group()`. | +| `pg.set_transfer_engine(engine)` | Reuse an external `TransferEngine`. | The engine must outlive all process groups. | +| `pg.get_active_ranks(backend)` | Return the backend active-rank tensor. | Used by EP fallback and recovery paths. | +| `pg.get_num_synced_ranks(backend)` | Return the number of locally activatable group slots. | Diagnostic helper. | +| `pg.get_peer_state(backend, ranks)` | Read locally mirrored activation readiness. | A lightweight, communication-free query. | +| `pg.activate_ranks(backend, ranks)` | Propose activation through the Coordinator. | A single call from any online rank is sufficient. | +| `pg.recover_ranks(backend, ranks)` | Propose activation through the Coordinator. | Compatibility alias to `activate_ranks`. | +| `pg.deactivate_ranks(backend, ranks)` | Propose deactivation through the Coordinator. | A single call from any online rank is sufficient. | +| `pg.join_group(backend)` | Confirm readiness for activation and remain blocked until activation actually occurs. | Used for scale-up, replacement, and in-place rejoin. | +| `pg.sync_after_failure(backend)` | Report current link observations, wait for reconciliation, and apply the latest group view. | Called automatically when `auto_sync_on_failure=True`; it may also be called manually. | + +### Supported distributed operations + +Mooncake PG implements the following `torch.distributed` APIs. Support may +depend on device type, dtype, PyTorch version, and whether the current backend is +`mooncake` or `mooncake-cpu`; run the PG tests on the target environment before +production use. + +| API family | Examples | Notes | +| --- | --- | --- | +| Collectives | `all_reduce`, `broadcast`, `all_gather`, `all_gather_into_tensor`, `reduce_scatter_tensor`, `all_to_all`, `barrier`, `reduce`, `gather`, `scatter` | Active ranks participate; inactive ranks are skipped by backend internals. | +| Async work | `dist.all_reduce(..., async_op=True)` | Wait on the returned work object, then synchronize the device stream as needed. | +| P2P | `isend`, `irecv`, `batch_isend_irecv` | Single-tensor P2P is routed through the Mooncake backend shim. | + +## Elastic recovery protocol + +Mooncake PG separates join preparation from membership activation. A joining or +recovering rank completes local warmup, calls `join_group()`, and waits. An +existing rank may poll local readiness and then issue the activation proposal. +The Coordinator validates and distributes the resulting membership. + +### Healthy-rank side + +```python +from mooncake import pg + +dist.init_process_group( + backend="mooncake", + rank=rank, + world_size=2, + pg_options=pg.MooncakeBackendOptions( + 3, # max_group_size + False, # is_extension + ), +) + +backend = dist.group.WORLD +join_ranks = [2] + +while not all(pg.get_peer_state(backend, join_ranks)): + # Continue serving, back off, or poll according to your scheduler policy. + pass + +pg.recover_ranks(backend, join_ranks) +``` + +### Joining-rank side + +```python +from mooncake import pg + +dist.init_process_group( + backend="mooncake", + rank=2, + world_size=3, + pg_options=pg.MooncakeBackendOptions( + 3, # max_group_size + True, # is_extension + ), +) + +backend = dist.group.WORLD + +# Collectives are local-only before join_group. Use this +# window for framework-specific preparation, for example: +# capture_cuda_graphs() +# warm_up_model() + +pg.join_group(backend) +``` + +Important semantics: + +- `get_peer_state()` is a local best-effort readiness query, not a collective. +- Capacity must be reserved with `max_group_size` when founding members create + the group. A joining registration appends inactive slots within that capacity. +- A joining rank starts with local-only collective behavior until `join_group`. + The join call then blocks until a Coordinator-approved activation commits. +- A single `activate_ranks()` call, or its `recover_ranks()` alias, from any + online rank is sufficient; redundant equivalent calls are safe. +- Subgroups must be created in the same order on healthy and joining processes, + following PyTorch `new_group()` ordering rules. + +## Mooncake EP quick start + +Mooncake EP exposes `Buffer` from `mooncake.mooncake_ep_buffer`. Initialize it +with a Mooncake process group and a workspace size computed from the expected +dispatch shape. + +```python +import torch +import torch.distributed as dist +from mooncake import pg +from mooncake.mooncake_ep_buffer import Buffer + + +# Assume dist.init_process_group(..., backend="mooncake", ...) has completed. +group = dist.group.WORLD +rank = dist.get_rank(group) +world_size = dist.get_world_size(group) + +num_tokens = 128 +hidden = 7168 +num_experts = 288 +top_k = 8 +max_tokens_per_rank = 128 + +x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device="cuda") +scores = torch.randn(num_tokens, num_experts, dtype=torch.float32, device="cuda") +topk_idx = torch.topk(scores, top_k, dim=-1).indices +topk_weights = torch.softmax( + torch.randn(num_tokens, top_k, dtype=torch.float32, device="cuda"), dim=-1 +) + +num_ep_buffer_bytes = Buffer.get_ep_buffer_size_hint( + max_tokens_per_rank, + hidden, + world_size, + num_experts, +) +buffer = Buffer(group, num_ep_buffer_bytes) + +# EP-level rank-health tensor. Kernels may update it to 0 when timeout_us +# detects a failed source rank. +active_ranks = torch.ones(world_size, dtype=torch.int32, device="cuda") + +recv_x, recv_count, handle, event, hook = buffer.dispatch( + x, + topk_idx, + active_ranks, + num_max_dispatch_tokens_per_rank=max_tokens_per_rank, + num_experts=num_experts, + timeout_us=-1, + use_fp8=True, + async_finish=False, + return_recv_hook=False, +) +event.current_stream_wait() + +# Run local experts on recv_x here. If use_fp8=True, recv_x is a +# (data, scales) tuple; dequantize or feed it into an FP8-aware expert kernel. +expert_out = run_local_experts(recv_x, recv_count) + +combined_x, event, hook = buffer.combine( + expert_out, + topk_idx, + topk_weights, + active_ranks, + timeout_us=-1, + handle=handle, +) +event.current_stream_wait() +``` + +## Mooncake EP API reference + +### `Buffer.get_ep_buffer_size_hint(...)` + +```python +Buffer.get_ep_buffer_size_hint( + num_max_dispatch_tokens_per_rank: int, + hidden: int, + num_ranks: int, + num_experts: int, +) -> int +``` + +Returns the workspace size in bytes for the EP buffer. Use the maximum number of +tokens a rank may dispatch in one step. Underestimating this value can cause +buffer overflow or incorrect dispatch results. + +### `Buffer(group, num_ep_buffer_bytes=0)` + +Creates the EP runtime for a Mooncake process group. The constructor exchanges +RDMA and IPC metadata through the group, initializes fast-path transports when +available, and falls back to the Python implementation if the fast path is not +usable. + +### `Buffer.dispatch(...)` + +```python +recv_x, recv_count, handle, event, hook = buffer.dispatch( + x, + topk_idx, + active_ranks, + num_max_dispatch_tokens_per_rank, + num_experts, + timeout_us, + use_fp8=True, + async_finish=False, + return_recv_hook=False, +) +``` + +Arguments: + +- `x`: local token hidden states, shape `[num_tokens, hidden]`, typically BF16 + on CUDA. +- `topk_idx`: selected expert IDs, shape `[num_tokens, top_k]`. Use `-1` to mark + masked selections. +- `active_ranks`: EP-level rank-health tensor, shape `[num_ranks]`, dtype + `torch.int32`. Timeout detection may set failed source ranks to `0`. +- `num_max_dispatch_tokens_per_rank`: workspace capacity per rank. It should be + at least the maximum local `num_tokens` across ranks for the current step. +- `num_experts`: global expert count. It must be divisible by `num_ranks`. +- `timeout_us`: timeout in microseconds. Use `-1` to disable timeout detection. +- `use_fp8`: when `True`, dispatch returns FP8 data plus scales. +- `async_finish`: when `True`, returned tensors are associated with the returned + event for stream-lifetime management. +- `return_recv_hook`: when `True`, call the returned `hook()` to complete receive + synchronization; otherwise use `event.current_stream_wait()`. + +Returns: + +- `recv_x`: packed local-expert inputs. If `use_fp8=True`, this is + `(packed_data, packed_scales)`; otherwise it is a BF16 tensor. +- `recv_count`: number of tokens received by each local expert. +- `handle`: opaque metadata required by `combine()` and + `get_next_combine_buffer()`. +- `event`: `EventOverlap` helper; call `event.current_stream_wait()` before using + outputs when no hook is used. +- `hook`: optional synchronization hook used when `return_recv_hook=True`. + +### `Buffer.combine(...)` + +```python +combined_x, event, hook = buffer.combine( + x, + topk_idx, + topk_weights, + active_ranks, + timeout_us, + handle, + zero_copy=False, + async_finish=False, + return_recv_hook=False, + out=None, +) +``` + +Arguments: + +- `x`: local expert outputs packed in the layout returned by `dispatch()`. +- `topk_idx` and `topk_weights`: routing metadata for combining expert outputs + back to local tokens. +- `active_ranks`: same EP-level rank-health tensor used by `dispatch()`. +- `timeout_us`: timeout in microseconds; use `-1` to disable timeout detection. +- `handle`: the handle returned by the matching `dispatch()` call. +- `zero_copy`: when `True`, write expert outputs into + `buffer.get_next_combine_buffer(handle)` and pass that tensor to `combine()`. +- `out`: optional output tensor for the combined result. + +### `Buffer.get_next_combine_buffer(handle)` + +Returns the next combine buffer for zero-copy expert output. Use it only with the +matching dispatch `handle` and pass the resulting tensor back to `combine()` with +`zero_copy=True`. + +### `Buffer.update_ep_member()` + +Reconnects EP peers after backend membership changes. Call it after PG recovery +updates rank activeness so EP transport metadata and QPs can be refreshed. + +## Active-rank tensors: PG vs EP + +There are two active-rank tensors in the API surface: + +- **PG active-rank mask**: passed to `pg.MooncakeBackendOptions`. This mirrors + the Coordinator's committed membership. +- **EP active-rank tensor**: passed to `Buffer.dispatch()` and `Buffer.combine()`. + It is also rank-level (`[num_ranks]`, `torch.int32`) and may be updated by EP + kernels when timeout detection marks a peer as failed. + +Their values may coincide in a simple integration, but their semantics are not +interchangeable: PG membership is configuration, while EP may update its mask +from kernel-level timeout observations. Keep the mapping, dtype, device, and +capacity consistent when propagating committed PG membership into EP. + +## Tests and examples + +- PG collectives: `mooncake-pg/tests/test_pg_collectives.py` +- PG elastic recovery and subgroup extension: `mooncake-pg/tests/test_pg_elastic.py` +- PG benchmark harness: `mooncake-pg/benchmark/README.md` +- EP correctness and failure simulation: `python/tests/ep/test_ep_grid.py` +- EP wrapper example: `python/tests/ep/test_mooncake_ep.py` + +See [PG/EP troubleshooting](../../troubleshooting/pg-ep-troubleshooting.md) for +common setup and runtime issues. diff --git a/docs/source/api-reference/python/index.md b/docs/source/api-reference/python/index.md new file mode 100644 index 0000000000..4c61577c59 --- /dev/null +++ b/docs/source/api-reference/python/index.md @@ -0,0 +1,18 @@ +# Python + +| Module | Description | +|--------|-------------| +| [Mooncake Store](mooncake-store) | Distributed KV cache storage client — `put`/`get`/`remove`/`replicate` operations | +| [DataProto Structured Object Transfer](dataproto-structured-object-transfer) | Structured-object helpers for storing and retrieving DataProto-like payloads | +| [Transfer Engine](transfer-engine) | High-performance RDMA/TCP data transfer between nodes | +| [EP Backend](ep-backend) | Expert-parallel backend for large MoE model deployment | + +:::{toctree} +:maxdepth: 1 +:hidden: + +mooncake-store +dataproto-structured-object-transfer +transfer-engine +ep-backend +::: diff --git a/docs/source/api-reference/python/mooncake-store.md b/docs/source/api-reference/python/mooncake-store.md new file mode 100644 index 0000000000..6a9f3cc9b4 --- /dev/null +++ b/docs/source/api-reference/python/mooncake-store.md @@ -0,0 +1,3283 @@ +# Mooncake Store Python API + +## Installation + +### PyPI Package +Install the Mooncake Transfer Engine package from PyPI, which includes both Mooncake Transfer Engine and Mooncake Store Python bindings: + +```bash +pip install mooncake-transfer-engine +``` + +📦 **Package Details**: [https://pypi.org/project/mooncake-transfer-engine/](https://pypi.org/project/mooncake-transfer-engine/) + +### Required Service +The only always-required service is: + +- `mooncake_master` — Master service for cluster membership and object placement + +For Transfer Engine metadata, use the `P2PHANDSHAKE` connection string for +decentralized peer discovery, enable the master's embedded HTTP metadata +server, or provide an external metadata service. + +## Quick Start + +### Start Master (with HTTP enabled) + +Enable the built-in HTTP metadata server when starting the master: + +```bash +mooncake_master \ + --enable_http_metadata_server=true \ + --http_metadata_server_host=0.0.0.0 \ + --http_metadata_server_port=8080 +``` +This exposes the metadata endpoint at `http://:/metadata`. + +### Hello World Example + +```python +from mooncake.store import MooncakeDistributedStore + +# 1. Create store instance +store = MooncakeDistributedStore() + +# 2. Setup with all required parameters +store.setup( + "localhost", # Your node's address + "http://localhost:8080/metadata", # HTTP metadata server + 512*1024*1024, # 512MB segment size + 128*1024*1024, # 128MB local buffer + "tcp", # Use TCP (RDMA for high performance) + "", # Leave empty; Mooncake auto-picks RDMA devices when needed + "localhost:50051" # Master service +) + +# 3. Store data +store.put("hello_key", b"Hello, Mooncake Store!") + +# 4. Retrieve data +data = store.get("hello_key") +print(data.decode()) # Output: Hello, Mooncake Store! + +# 5. Clean up +store.close() +``` + +**RDMA device selection**: For `protocol="rdma"` or `protocol="efa"`, leave +`rdma_devices` as `""` to auto-discover NICs. Set `MC_MS_AUTO_DISC=0` when you +want auto-discovery disabled, then provide a comma-separated list such as +`"mlx5_0,mlx5_1"` to pin specific hardware. + +Mooncake selects available ports internally at `setup() `, so you do not need to fix specific port numbers in these examples. Internally, ports are chosen from a dynamic range (currently 12300–14300). + +#### P2P Hello World + +The following setup uses P2P handshake and does not require an HTTP metadata +server. Pass the literal `P2PHANDSHAKE` value as the metadata server. + +```python +from mooncake.store import MooncakeDistributedStore + +store = MooncakeDistributedStore() +store.setup( + "localhost", # Your node's ip address + "P2PHANDSHAKE", # P2P handshake (no HTTP metadata) + 512*1024*1024, # 512MB segment size + 128*1024*1024, # 128MB local buffer + "tcp", # Use TCP (RDMA for high performance) + "", # Leave empty; Mooncake auto-picks RDMA devices when needed + "localhost:50051" # Master service +) + +store.put("hello_key", b"Hello, Mooncake Store!") +print(store.get("hello_key").decode()) +store.close() +``` + +## Basic API Usage + +### Simple Get/Put Operations + +
+Click to expand: Complete Get/Put example with NumPy arrays + +```python +import numpy as np +import json +from mooncake.store import MooncakeDistributedStore + +# 1. Initialize +store = MooncakeDistributedStore() +store.setup("localhost", + "http://localhost:8080/metadata", + 512*1024*1024, + 128*1024*1024, + "tcp", + "", + "localhost:50051") +print("Store ready.") + +# 2. Store data +store.put("config", b'{"model": "llama-7b", "temperature": 0.7}') +model_weights = np.random.randn(1000, 1000).astype(np.float32) +store.put("weights", model_weights.tobytes()) +store.put("cache", b"some serialized cache data") + +# 3. Retrieve and verify data +config = json.loads(store.get("config").decode()) +weights = np.frombuffer(store.get("weights"), dtype=np.float32).reshape(1000, 1000) + +print("Config OK:", config["model"]) +print("Weights OK, mean =", round(float(weights.mean()), 4)) +print("Cache exists?", bool(store.is_exist("cache"))) + +# 4. Close +store.close() +``` + +
+ +## Structured Object Store Helper + +`mooncake.structured_object_store` provides a higher-level helper for one logical object that contains multiple named members. +It is designed for cases such as rollout / batch transfer where callers want to keep their own object semantics locally while using Mooncake for fast payload movement. + +The helper separates two concepts: + +- **structured object path**: named members with metadata-aware materialization; +- **generic bundle path**: manifest + named payloads when the caller only needs raw grouped objects. + +### Main types + +```python +from mooncake.structured_object_store import ( + MooncakeBundleTransfer, + StructuredMemberSlice, + StructuredObjectPayload, +) +``` + +- `MooncakeBundleTransfer`: public helper facade built on a `MooncakeDistributedStore`. +- `StructuredObjectPayload`: structured object to write. Members are passed in `buffers`, and optional object metadata is passed in `metadata`. +- `StructuredMemberSlice`: slice selection for one structured member during reads. + +### Structured object write and full read + +Use `put_structured_object()` to write one structured object. The default read path is `read_spec(ref)`, and full-object materialization is just the default case of the partial-read API. + +```python +import numpy as np +from mooncake.store import MooncakeDistributedStore +from mooncake.structured_object_store import MooncakeBundleTransfer, StructuredObjectPayload + +store = MooncakeDistributedStore() +transfer = MooncakeBundleTransfer(store, key_prefix="demo/structured") + +payload = StructuredObjectPayload( + metadata={"step": 7, "layout": "rollout"}, + buffers={ + "tokens": np.array(range(24), dtype=np.int32).reshape(6, 4), + "mask": np.ones((6, 4), dtype=np.int8), + "prompt_ids": b"sample-ids", + }, +) + +ref = transfer.put_structured_object(payload) +result = transfer.materialize(transfer.read_spec(ref)) + +tokens = result.objects["tokens"] +prompt_ids = result.objects["prompt_ids"] +metadata = result.metadata +``` + +### Structured object transfer policy + +By default, structured object writes use `BundleTransferPolicy(copy_mode="auto")`. Non-tensor payloads are staged through a Mooncake BufferPool and written with `batch_put_from` when that path is available; otherwise Mooncake falls back to ordinary `store.put`. Typed-ragged ndarray rows use the native fast-copy extension to copy directly into the staging buffer without an intermediate concatenation. + +Use `copy_mode="copy"` to force the regular `store.put` path: + +```python +from mooncake.structured_object_store import BundleTransferPolicy + +ref = transfer.put_structured_object( + payload, + policy=BundleTransferPolicy(copy_mode="copy"), +) +``` + +Available modes: + +- `auto`: prefer BufferPool staging plus `batch_put_from` for non-tensor payloads and fall back to regular `store.put` when that path is unavailable; +- `copy`: force the regular `store.put` path; +- `zero_copy`: require payloads to be explicit `tensor_object_buffer` instances; non-tensor structured payloads are rejected. + +### Partial reads + +Read narrowing happens on top of `read_spec(ref)`: + +- `select_members([...])` keeps only selected members; +- `slice_member(name, axis=0, start=..., end=...)` slices one ndarray member; +- `materialize(spec)` returns newly materialized objects. + +```python +spec = ( + transfer.read_spec(ref) + .select_members(["tokens"]) + .slice_member("tokens", axis=0, start=2, end=5) +) +result = transfer.materialize(spec) + +selected_tokens = result.objects["tokens"] +``` + +Current scope: + +- byte members support full-member reads; +- ndarray members support full reads and sliced reads; +- full read is the default `read_spec(ref)` case. + +### Reusing caller-owned destinations + +Use `materialize_into()` when the caller already owns the destination ndarray buffers and wants Mooncake to fill them directly. + +```python +destination = np.empty((3, 4), dtype=np.int32) +spec = ( + transfer.read_spec(ref) + .select_members(["tokens"]) + .slice_member("tokens", axis=0, start=2, end=5) +) +result = transfer.materialize_into(spec, {"tokens": destination}) + +assert result.objects["tokens"] is destination +``` + +`materialize_into()` is only for members whose destination layout is already known to the caller. For byte members or default object reconstruction, use `materialize()`. + +### Generic bundle fallback + +If the caller does not need structured member semantics, the same helper also supports raw named bundles: + +- `put_bundle(...)` +- `remove_bundle(...)` + +Use the bundle path when the object is just a manifest plus named payloads, and use the structured object path when callers want member selection, slicing, and ndarray-aware materialization. + +## Zero-Copy API (Advanced Performance) + +For maximum performance, especially with RDMA networks, use the zero-copy API. This allows direct memory access without intermediate copies. + +### Memory Registration + +⚠️ **Important**: `register_buffer` is required for zero-copy RDMA operations. Without proper buffer registration, undefined behavior and memory corruption may occur. + +Zero-copy operations require registered memory buffers. For repeated reads and writes, prefer the Python `BufferPool` helper described below so leases come from the store's setup-time local buffer instead of registering and unregistering memory for every operation. + +#### register_buffer() +Register a memory buffer for direct RDMA access. + +#### unregister_buffer() +Unregister a previously registered buffer. + +
+Click to expand: Buffer registration example + +```python +import numpy as np +from mooncake.store import MooncakeDistributedStore + +# Initialize store +store = MooncakeDistributedStore() +store.setup("localhost", "http://localhost:8080/metadata", 512*1024*1024, 128*1024*1024, "tcp", "", "localhost:50051") + +# Create a large buffer +buffer = np.zeros(100 * 1024 * 1024, dtype=np.uint8) # 100MB buffer + +# Register the buffer for zero-copy operations +buffer_ptr = buffer.ctypes.data +result = store.register_buffer(buffer_ptr, buffer.nbytes) +if result != 0: + print(f"Failed to register buffer: {result}") + raise RuntimeError(f"Failed to register buffer: {result}") +print("Buffer registered successfully.") +store.unregister_buffer(buffer_ptr) +``` + +
+ +#### BufferPool Helper + +`BufferPool` leases scratch buffers from the store's setup-time local buffer for repeated zero-copy operations. It is useful when a caller repeatedly needs temporary memory, for example as the destination buffer for `get_into()` or `get_into_ranges()`. + +```python +from mooncake.buffer_pool import BufferPool + +pool = BufferPool(store) + +with pool.buffer(1024 * 1024) as lease: + n = store.get_into("my_key", lease.ptr, lease.size) + view = lease.buffer[:n] + # Consume view directly, or wrap it with np.frombuffer(view, dtype=...). + # Copy only if the data must outlive the lease: data = bytes(view) + +pool.close() +``` + +`acquire(size)` and `buffer(size)` return a lease object. A lease exposes: + +- `ptr`: the local-buffer address to pass to zero-copy APIs. +- `size`: the requested logical size. +- `buffer`: a Python `memoryview` over the logical requested size. +- `release()`: returns the local-buffer allocation to the store allocator. + +Behavior and lifecycle rules: + +- Leases prefer the store local buffer shared with internal Store staging paths. +- If local-buffer allocation is temporarily exhausted, `BufferPool` can allocate and register a short-lived overflow buffer; the overflow buffer is unregistered when the lease is released. +- `max_regions` can limit the number of concurrently active external leases. +- `max_bytes` bounds total active local-buffer and overflow leases; the default allows one local-buffer-sized overflow burst. +- `acquire(size, block=False)` raises when both local and overflow capacity are exhausted instead of waiting. +- `acquire(size, timeout=...)` can wait for another lease to be released. +- Do not keep `lease.ptr` or a `memoryview` after releasing the lease. +- `release()` fails while exported views are alive. Delete those views first, then release. +- `close()` fails if leases are still active. Release all leases before closing the pool. + +Legacy code may still import `RegisteredBufferPool`, but new examples should prefer `mooncake.buffer_pool.BufferPool`. + +--- + +### Zero-Copy Operations + +#### Complete Zero-Copy Workflow + +⚠️ **Critical**: Always register buffers before zero-copy operations. Failure to register buffers will cause undefined behavior and potential memory corruption. + +Here's a complete example showing the full zero-copy workflow with proper buffer management: + +
+Click to expand: Complete zero-copy workflow example + +```python +import numpy as np +from mooncake.store import MooncakeDistributedStore + +# Initialize store with RDMA protocol for maximum performance +store = MooncakeDistributedStore() +store.setup("localhost", "http://localhost:8080/metadata", 512*1024*1024, 16*1024*1024, "tcp", "", "localhost:50051") + +# Create data to store +original_data = np.random.randn(1000, 1000).astype(np.float32) +buffer_ptr = original_data.ctypes.data +size = original_data.nbytes + +# Step 1: Register the buffer +result = store.register_buffer(buffer_ptr, size) +if result != 0: + raise RuntimeError(f"Failed to register buffer: {result}") + +# Step 2: Zero-copy store +result = store.put_from("large_tensor", buffer_ptr, size) +if result == 0: + print(f"Successfully stored {size} bytes with zero-copy") +else: + raise RuntimeError(f"Store failed with code: {result}") + +# Step 3: Pre-allocate buffer for retrieval +retrieved_data = np.empty((1000, 1000), dtype=np.float32) +recv_buffer_ptr = retrieved_data.ctypes.data +recv_size = retrieved_data.nbytes + +# Step 4: Register receive buffer +result = store.register_buffer(recv_buffer_ptr, recv_size) +if result != 0: + raise RuntimeError(f"Failed to register receive buffer: {result}") + +# Step 5: Zero-copy retrieval +bytes_read = store.get_into("large_tensor", recv_buffer_ptr, recv_size) +if bytes_read > 0: + print(f"Successfully retrieved {bytes_read} bytes with zero-copy") + # Verify the data + print(f"Data matches: {np.array_equal(original_data, retrieved_data)}") +else: + raise RuntimeError(f"Retrieval failed with code: {bytes_read}") + +# Step 6: Clean up - unregister both buffers +store.unregister_buffer(buffer_ptr) +store.unregister_buffer(recv_buffer_ptr) +store.close() +``` + +
+ +#### put_from() +Store data directly from a registered buffer (zero-copy). + +```python +def put_from(self, key: str, buffer_ptr: int, size: int, config=None) -> int +``` + +**Parameters:** +- `key`: Object identifier +- `buffer_ptr`: Memory address (from ctypes.data or similar) +- `size`: Number of bytes to store +- `config`: Optional replication configuration + +#### get_into() +Retrieve data directly into a registered buffer (zero-copy). + +```python +def get_into(self, key: str, buffer_ptr: int, size: int) -> int +``` + +**Parameters:** +- `key`: Object identifier to retrieve +- `buffer_ptr`: Memory address of pre-allocated buffer +- `size`: Size of the buffer (must be >= object size) + +**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. Each buffer must resolve to Store-managed registered memory, either from `BufferPool`/the setup-time local buffer or from an explicit `register_buffer()` call. +- `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 + +The `ReplicateConfig` class allows you to control data replication behavior when storing objects in Mooncake Store. This configuration is essential for ensuring data reliability, performance optimization, and storage placement control. + +### Class Definition + +```python +from mooncake.store import ReplicateConfig + +# Create a configuration instance +config = ReplicateConfig() +``` + +### Properties + +#### replica_num +**Type:** `int` +**Default:** `1` +**Description:** Specifies the number of memory replicas to create for the +stored object. + +```python +config = ReplicateConfig() +config.replica_num = 3 # Store 3 memory replicas +``` + +#### nof_replica_num +**Type:** `int` +**Default:** `0` +**Description:** Specifies the number of replicas to create in the configured +NVMe-oF SSD pool. + +```python +config = ReplicateConfig() +config.replica_num = 1 +config.nof_replica_num = 1 +``` + +#### dfs_replica_num +**Type:** `int` +**Default:** `0` +**Status:** **Work in progress; development and evaluation only.** +**Description:** Requests an additional replica in the configured shared +distributed filesystem. The supported values are currently `0` and `1`. When +set to `1`, `replica_num` must be at least `1`, so DFS-only placement is not +supported. DFS replicas currently support only the `default` tenant. + +```python +config = ReplicateConfig() +config.replica_num = 1 +config.dfs_replica_num = 1 +``` + +Writes that request a DFS replica return success after the DFS `WriteAt` +operation completes, but without an additional `fsync` durability guarantee. +The master and client DFS backends must be enabled and configured with the same +absolute shared-root path and shard layout. See the +{ref}`DFS deployment documentation ` for the required environment +variables and current limitations. + +For a same-size `upsert`, if either the existing object or the new request has +a DFS replica, the requested memory, NoF, and DFS replica counts must match the +existing topology. A different-size update allocates a new topology. + +#### soft_pin_action +**Type:** `SoftPinAction` +**Default:** `SoftPinAction.PRESERVE` +**Description:** Controls the soft-pin transition committed when the first replica becomes readable. `PRESERVE` keeps an existing deadline during Upsert, `ENABLE` starts a fixed soft-pin lifetime, and `DISABLE` removes it. Reads do not extend the lifetime. + +```python +from mooncake.store import ReplicateConfig, SoftPinAction + +config = ReplicateConfig() +config.soft_pin_action = SoftPinAction.ENABLE +config.soft_pin_ttl_ms = 60_000 # Optional; omitted uses the Master default +``` + +`soft_pin_ttl_ms` is valid only with `ENABLE`. The Master rejects TTLs above +`max_kv_soft_pin_ttl`; a value of zero commits the object as ordinary cache. +Soft-pin state is not persisted in snapshots or the HA OpLog. Restored objects +therefore become ordinary cache after recovery or Standby promotion. + +#### 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) +**Description:** Specifies a preferred segment (node) for data allocation. This is typically the hostname:port of a target server. + +```python +config = ReplicateConfig() + +# Preferred replica location ("host:port") +config.preferred_segment = "localhost:12345" # pin to a specific machine + +# Alternatively, pin to the local host +config.preferred_segment = self.get_hostname() + +# Optional: speed up local transfers +# export MC_STORE_MEMCPY=1 +``` + +#### prefer_alloc_in_same_node +**Type:** `bool` +**Default:** `False` +**Description:** Enables host-aware local-first allocation for this request, using the writer host identity and the normal ordered remote fallback. This can be used with direct multi-buffer writes and tensor write APIs to avoid staging when the selected segment is local and local memcpy is enabled. It does not support disk segments, and the `replica_num` can only be set to 1. Tensor APIs keep their default staging behavior unless this flag is explicitly enabled. + +```python +config = ReplicateConfig() +config.prefer_alloc_in_same_node = True +``` + +#### group_ids +**Type:** `List[str] | None` +**Default:** `None` +**Description:** Optionally assigns object metadata to routing groups during writes. When this field is unset, Mooncake Store preserves the default ungrouped behavior. When it is set, each group ID maps to the object at the same position in the write request. Empty string (`""`) explicitly stores that object as ungrouped. + +For batch write APIs, the number of group IDs must match the number of keys: + +```python +config = ReplicateConfig() +config.group_ids = ["session-a", "", "session-b"] + +store.put_batch( + ["key-a", "key-b", "key-c"], + [b"value-a", b"value-b", b"value-c"], + config, +) +``` + +For a single-object write, provide one group ID: + +```python +config = ReplicateConfig() +config.group_ids = ["session-a"] + +store.put("key-a", b"value-a", config) +``` + +--- + +## Unified Parallel Tensor IO API + +Mooncake Store also provides a unified tensor IO family for tensors that are stored either as full objects or as explicitly identified parallel shards. + +This API family is the long-term interface for TP / DP / EP / PP-aware tensor IO: + +- write and upsert use `TensorParallelism` +- reads use `ReadTarget` +- legacy TP-only APIs remain available as compatibility wrappers + +### ParallelAxis + +`ParallelAxis` describes one axis in a shard identity. + +```python +axis = mooncake.store.ParallelAxis() +axis.kind = "tp" # one of: "tp", "dp", "ep", "pp" +axis.rank = 0 +axis.size = 8 +axis.split_dim = 1 # used for layout-sharding axes such as TP +axis.expert_id = 3 # optional, for EP +axis.stage_id = 1 # optional, for PP +``` + +**Fields:** +- `kind`: Parallelism axis kind. +- `rank`: Current shard rank on that axis. +- `size`: Total number of shards on that axis. +- `split_dim`: Optional tensor split dimension for layout-sharding axes. +- `expert_id`: Optional expert identifier for EP layouts. +- `stage_id`: Optional pipeline stage identifier for PP layouts. + +### TensorParallelism + +`TensorParallelism` is an ordered list of axes that identifies the stored or requested shard. + +```python +parallelism = mooncake.store.TensorParallelism() +parallelism.axes = [ + tp_axis, +] +``` + +Examples: +- TP shard: `axes=[TP(...)]` +- DP + TP shard: `axes=[DP(...), TP(...)]` +- PP + TP shard: `axes=[PP(...), TP(...)]` +- EP shard: `axes=[EP(...)]` + +### ReadTarget + +`ReadTarget` tells Mooncake whether the caller wants the stored form, a specific shard view, or the reconstructed full tensor. + +```python +target = mooncake.store.ReadTarget() +target.mode = "full" # one of: "as_stored", "shard", "full" +target.parallelism = None # required for target shard reads +``` + +**Fields:** +- `mode`: Read materialization mode. +- `parallelism`: Optional `TensorParallelism`. Required when `mode="shard"`. + +### put_tensor_with_parallelism() + +Store a tensor using the unified parallelism model. + +```python +def put_tensor_with_parallelism( + self, + key: str, + tensor, + parallelism: mooncake.store.TensorParallelism | None = None, + config: ReplicateConfig | None = None, + writer_partition = None, +) -> int +``` + +Use `parallelism=None` to store a full tensor object. Provide `TensorParallelism` to store a shard-scoped object. + +`writer_partition` is an optional write-side shorthand for full-tensor inputs that should be stored as one shard. It describes the writer's `(rank, size, split_dim)` and is mutually exclusive with `parallelism`; do not provide both in one call. + +For TP-containing multi-axis layouts, the caller may pass the full source tensor; Mooncake derives and persists the uniform shard selected by the requested TP rank/layout. That applies to layouts such as `dp_tp`, `pp_tp`, and `ep_tp`. + +Plain single-axis TP remains shard-input for compatibility. + +Pure DP still does not imply a split axis by itself. + +### batch_put_tensor_with_parallelism() + +Batch version of unified tensor writes. + +```python +def batch_put_tensor_with_parallelism( + self, + keys: list[str], + tensors: list, + parallelisms: list[mooncake.store.TensorParallelism | None] | None = None, + config: ReplicateConfig | None = None, + writer_partitions = None, +) -> list[int] +``` + +`writer_partitions` is an optional write-side convenience input for batch full-tensor writes that should be partitioned into stored shards. Each entry describes the target shard write as `(rank, size, split_dim)`. + +Use `writer_partitions` when the caller has full tensors and wants Mooncake to derive the stored shard objects from writer-side partition info instead of constructing full `TensorParallelism` objects per element. TP-containing `parallelisms` can now express the same full-tensor-input behavior too; `writer_partitions` remains the lighter explicit write-side shorthand. + +### get_tensor_with_parallelism() + +Read a tensor through the unified read path. + +```python +def get_tensor_with_parallelism( + self, + key: str, + target: mooncake.store.ReadTarget | None = None, +) +``` + +Typical modes: +- `target=None` or `mode="as_stored"`: return the stored local object. +- `mode="shard"`: return the target shard described by `target.parallelism`. +- `mode="full"`: reconstruct and return the full tensor. + +### batch_get_tensor_with_parallelism() + +Batch version of unified tensor reads. + +```python +def batch_get_tensor_with_parallelism( + self, + keys: list[str], + targets: list[mooncake.store.ReadTarget | None] | None = None, +) -> list +``` + +### get_tensor_with_parallelism_into() / batch_get_tensor_with_parallelism_into() + +Zero-copy unified read forms. The destination buffers must resolve to Store-managed registered memory, either from `BufferPool`/the setup-time local buffer or from an explicit `register_buffer()` call. + +```python +def get_tensor_with_parallelism_into( + self, + key: str, + buffer_ptr: int, + size: int, + target: mooncake.store.ReadTarget | None = None, +) +``` + +```python +def batch_get_tensor_with_parallelism_into( + self, + keys: list[str], + buffer_ptrs: list[int], + sizes: list[int], + targets: list[mooncake.store.ReadTarget | None] | None = None, +) -> list +``` + +### upsert_tensor_with_parallelism() + +Unified upsert form for tensor objects. + +```python +def upsert_tensor_with_parallelism( + self, + key: str, + tensor, + parallelism: mooncake.store.TensorParallelism | None = None, + config: ReplicateConfig | None = None, + writer_partition = None, +) -> int +``` + +The write semantics match `put_tensor_with_parallelism()`, including full-tensor input for TP-containing layouts and the mutually exclusive `writer_partition` shorthand. + +### batch_upsert_tensor_with_parallelism() + +Batch unified upsert form. + +```python +def batch_upsert_tensor_with_parallelism( + self, + keys: list[str], + tensors: list, + parallelisms: list[mooncake.store.TensorParallelism | None] | None = None, + config: ReplicateConfig | None = None, + writer_partitions = None, +) -> list[int] +``` + +The write semantics match `put_tensor_with_parallelism()`, including full-tensor input for TP-containing layouts. + +### *_from zero-copy write variants + +The unified write and upsert family also has `_from` variants for registered-memory inputs, including: + +- `put_tensor_with_parallelism_from(...)` +- `batch_put_tensor_with_parallelism_from(...)` +- `upsert_tensor_with_parallelism_from(...)` +- `batch_upsert_tensor_with_parallelism_from(...)` + +These APIs accept Store-managed registered buffer pointers that contain serialized tensor objects in the current Mooncake tensor format: + +```text +[TensorObjectHeader + layout metadata][tensor data] +``` + +As with other zero-copy APIs, every source pointer must resolve to Store-managed registered memory, either from `BufferPool`/the setup-time local buffer or from an explicit `register_buffer()` call. + +### Compatibility wrappers + +Legacy TP-only methods such as: + +- `put_tensor_with_tp(...)` +- `batch_put_tensor_with_tp(...)` +- `get_tensor_with_tp(...)` +- `batch_get_tensor_with_tp(...)` +- corresponding `_into`, `_from`, and upsert variants + +remain supported for compatibility, but they are wrapper-style APIs around the unified parallel tensor IO model. Prefer the unified `*_with_parallelism` family for new code and new documentation examples. + +--- + +## Non-Zero-Copy API (Simple Usage) + +For simpler use cases, use the standard API without memory registration: + +### Basic Operations + +
+Click to expand: Non-zero-copy API examples + +```python +from mooncake.store import MooncakeDistributedStore + +# Initialize (same as zero-copy) +store = MooncakeDistributedStore() +store.setup("localhost", "http://localhost:8080/metadata", 512*1024*1024, 128*1024*1024, "tcp", "", "localhost:50051") + +# Simple put/get (automatic memory management) +data = b"Hello, World!" * 1000 # ~13KB +store.put("message", data) + +retrieved = store.get("message") +print(retrieved == data) # True + +# Batch operations +keys = ["key1", "key2", "key3"] +values = [b"value1", b"value2", b"value3"] + +store.put_batch(keys, values) +retrieved = store.get_batch(keys) +print("Retrieved all keys successfully:", retrieved == values) +``` + +
+ +### Performance Notes + +**Choose the appropriate API based on your use case:** + +**Zero-copy API is beneficial when:** +- Working with large data transfers +- RDMA network infrastructure is available and configured +- Direct memory access patterns fit your application design + +**Non-zero-copy API is suitable for:** +- Development and prototyping phases +- Applications without specific performance requirements + +**Batch operations can improve throughput for:** +- Multiple related operations performed together +- Scenarios where network round-trip reduction is beneficial + +--- + +## Topology & Devices + +- Auto-discovery: Enabled by default for `protocol="rdma"` or `protocol="efa"` + when `rdma_devices` is empty. +- Discovery controls: + - `MC_MS_AUTO_DISC=1` forces auto-discovery; then `rdma_devices` is ignored. + - `MC_MS_AUTO_DISC=0` disables auto-discovery; then `rdma_devices` is required + for RDMA/EFA. + - `MC_MS_FILTERS` restricts auto-discovery to a comma-separated whitelist of + NIC names, e.g. `MC_MS_FILTERS=mlx5_0,mlx5_2`. + +Examples: + +```bash +# Auto-select with default settings +python - <<'PY' +from mooncake.store import MooncakeDistributedStore as S +s = S() +s.setup("localhost", "http://localhost:8080/metadata", 512*1024*1024, 128*1024*1024, "rdma", "", "localhost:50051") +PY + +# Manual device list +unset MC_MS_AUTO_DISC +python - <<'PY' +from mooncake.store import MooncakeDistributedStore as S +s = S() +s.setup("localhost", "http://localhost:8080/metadata", 512*1024*1024, 128*1024*1024, "rdma", "mlx5_0,mlx5_1", "localhost:50051") +PY + +# Auto-select with filters +export MC_MS_AUTO_DISC=1 +export MC_MS_FILTERS=mlx5_0,mlx5_2 +python - <<'PY' +from mooncake.store import MooncakeDistributedStore as S +s = S() +s.setup("localhost", "http://localhost:8080/metadata", 512*1024*1024, 128*1024*1024, "rdma", "", "localhost:50051") +PY +``` + +## get_buffer Buffer Protocol + +The `get_buffer` method returns a `BufferHandle` object that implements the Python buffer protocol: + +
+Click to expand: Buffer protocol usage example + +```python +# Get buffer with buffer protocol support +buffer = store.get_buffer("large_object") +if buffer: + # Access as numpy array without copy + import numpy as np + arr = np.array(buffer, copy=False) + + # Direct memory access + ptr = buffer.ptr() # Memory address + size = buffer.size() # Buffer size in bytes + + # Use with other libraries that accept buffer protocol + print(f"Buffer size: {len(buffer)} bytes") + + # The buffer is automatically freed +``` + +
+--- + +## Full API Reference + +### Class: MooncakeDistributedStore + +The main class for interacting with Mooncake Store. + +#### Constructor +```python +store = MooncakeDistributedStore() +``` +Creates a new store instance. No parameters required. + +--- + + +#### setup() +Initialize distributed resources and establish network connections. + +```python +def setup( + self, + local_hostname: str, + metadata_server: str, + global_segment_size: int, + local_buffer_size: int, + protocol: str, + rdma_devices: str, + master_server_addr: str, + engine: Optional[TransferEngine] = None, + enable_ssd_offload: bool = False, + ssd_offload_path: str = "", + tenant_id: str = "default", + enable_client_http_server: bool = False, + client_http_port: int = 9300, +) -> int +``` + +The positional overload requires every argument through +`master_server_addr`. To use defaults for those fields, pass a configuration +dictionary instead: + +```python +def setup(self, config: Dict[str, object]) -> int +``` + +The dictionary overload requires `local_hostname` and `metadata_server`. Its +other keys are optional; the defaults are `16777216` (16 MiB) for both +`global_segment_size` and `local_buffer_size`, `"tcp"` for `protocol`, an empty +string for `rdma_devices`, and `"127.0.0.1:50051"` for +`master_server_addr`. It also accepts `ipc_socket_path` and the optional +configuration fields listed below. The `engine` argument is available only in +the positional overload. + +**Parameters:** +- `local_hostname` (str): **Required**. Local hostname and port (e.g., "localhost" or "localhost:12345") +- `metadata_server` (str): **Required**. Metadata connection string, e.g. `"P2PHANDSHAKE"` or `"http://localhost:8080/metadata"`. +- `global_segment_size` (int): **Required by the positional overload**. Memory segment size in bytes for mounting. +- `local_buffer_size` (int): **Required by the positional overload**. Local buffer size in bytes. +- `protocol` (str): **Required by the positional overload**. Network protocol, usually `"tcp"`, `"rdma"`, `"efa"`, `"cxl"`, or `"ascend"` depending on the build. +- `rdma_devices` (str): **Required by the positional overload**. RDMA/EFA device name(s), e.g. `"mlx5_0"` or `"mlx5_0,mlx5_1"`. Leave empty to auto-discover NICs unless `MC_MS_AUTO_DISC=0`; always empty for TCP. +- `master_server_addr` (str): **Required by the positional overload**. Master server address (e.g., "localhost:50051") +- `engine` (Optional[TransferEngine]): Existing Transfer Engine instance to reuse. Defaults to `None`. +- `enable_ssd_offload` (bool): Initialize client-side `FileStorage`. With a + normal file backend this enables SSD offload; with + `MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR=distributed_storage_backend`, + it initializes the DFS backend and is required for DFS reads and writes. + Defaults to `False`. +- `ssd_offload_path` (str): FileStorage directory. When provided, it overrides + `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH`. With the distributed backend, DFS shard + data is stored under `MOONCAKE_DFS_ROOT_DIR`, but this separate directory is + still validated during FileStorage initialization. +- `tenant_id` (str): Tenant namespace for object keys. Defaults to `"default"`. +- `enable_client_http_server` (bool): Enable the client-local `/health`, `/metrics`, `/metrics/summary`, and `/version` HTTP endpoints. Defaults to `False`. +- `client_http_port` (int): Port for the client-local HTTP endpoints. Defaults to `9300`. + +**Store segment pinned memory:** CUDA-enabled builds can register Store-managed +host segments as pinned memory when `MC_STORE_PIN_MEMORY_MAX_BYTES` is set to a +positive process-wide quota; unset, empty, `0`, or invalid values disable it. +The scope is limited to host Store segments allocated by `setup()` +(`global_segment_size`) and `allocateAndMountSegment()`; it excludes file-backed +`mountSegment()` mappings, CXL/device segments, `local_buffer_size`, user +buffers, dummy-client shared memory, and temporary staging buffers. If the quota +is exhausted or CUDA registration fails, Mooncake continues with pageable Store +segment memory. + +**Returns:** +- `int`: Status code (0 = success, non-zero = error code) + +**Example:** + +
+Click to expand: Setup examples for TCP and RDMA + +```python +# TCP initialization +store.setup("localhost", "http://localhost:8080/metadata", 1024*1024*1024, 128*1024*1024, "tcp", "", "localhost:50051") + +# RDMA auto-detect +store.setup("localhost", "http://localhost:8080/metadata", 512*1024*1024, 128*1024*1024, "rdma", "", "localhost:50051") + +# RDMA with explicit device list +store.setup("localhost", "http://localhost:8080/metadata", 512*1024*1024, 128*1024*1024, "rdma", "mlx5_0,mlx5_1", "localhost:50051") +``` + +
+ +--- +#### setup_dummy() +Initialize the store with a dummy client for testing purposes. + +```python +def setup_dummy(self, mem_pool_size: int, local_buffer_size: int, server_address: str) -> int +``` + +**Parameters:** +- `mem_pool_size` (int): Memory pool size in bytes +- `local_buffer_size` (int): Local buffer size in bytes +- `server_address` (str): Server address in format "hostname:port" + +**Returns:** +- `int`: Status code (0 = success, non-zero = error code) + +**Example:** +```python +# Initialize with dummy client +store.setup_dummy(1024*1024*256, 1024*1024*64, "localhost:8080") +``` + +Dummy clients do not own Store segments. They use a local shared-memory buffer +that is mapped by a real client process at `server_address`. Tensor APIs that +stage through this SHM buffer are supported, including tensor put/get, +`*_tensor_from`, `*_tensor_into`, tensor upsert/pub, TP wrappers, and unified +parallelism write wrappers. + +The real client owns the SHM buffer allocator. This keeps tensor writes and +regular object writes from allocating overlapping offsets when they run +concurrently through the same dummy client. Writes staged through the dummy +client's local SHM buffer keep their allocation alive until completion through +the real-side active buffer handle and dummy-side RAII release path. + +Full materialized reconstruction reads for writer-sharded or reconstructed +parallel tensors are still conservative for dummy clients. Use the corresponding +`*_into` APIs, or read stored shards directly, when using dummy clients. + +--- + +#### put() +Store binary data in the distributed storage. + +```python +def put(self, key: str, value: bytes, config: ReplicateConfig = None) -> int +``` + +**Parameters:** +- `key` (str): Unique object identifier +- `value` (bytes): Binary data to store +- `config` (ReplicateConfig, optional): Replication configuration + +**Returns:** +- `int`: Status code (0 = success, non-zero = error code) + +**Example:** + +
+Click to expand: Put operation examples + +```python +# Simple put +store.put("my_key", b"Hello, World!") + +# Put with replication config +config = ReplicateConfig() +config.replica_num = 2 +store.put("important_data", b"Critical information", config) +``` + +
+ +--- + +#### get() +Retrieve binary data from distributed storage. + +```python +def get(self, key: str) -> bytes +``` + +**Parameters:** +- `key` (str): Object identifier to retrieve + +**Returns:** +- `bytes`: Retrieved binary data + +**Raises:** +- Returns empty bytes if key doesn't exist + +**Example:** + +
+Click to expand: Get operation example + +```python +data = store.get("my_key") +if data: + print(f"Retrieved: {data.decode()}") +else: + print("Key not found") +``` + +
+ +--- + +#### put_batch() +Store multiple objects in a single batch operation. + +```python +def put_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 store +- `config` (ReplicateConfig, optional): Replication configuration for all objects + +**Returns:** +- `int`: Status code (0 = success, non-zero = error code) + +**Example:** + +
+Click to expand: Batch put example + +```python +keys = ["key1", "key2", "key3"] +values = [b"value1", b"value2", b"value3"] +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. + +```python +def get_batch(self, keys: List[str]) -> List[bytes] +``` + +**Parameters:** +- `keys` (List[str]): List of object identifiers to retrieve + +**Returns:** +- `List[bytes]`: List of retrieved binary data + +**Example:** + +
+Click to expand: Batch get example + +```python +keys = ["key1", "key2", "key3"] +values = store.get_batch(keys) +for key, value in zip(keys, values): + print(f"{key}: {len(value)} bytes") +``` + +
+ +--- + +#### remove() +Delete an object from the storage system. + +```python +def remove(self, key: str) -> int +``` + +**Parameters:** +- `key` (str): Object identifier to remove + +**Returns:** +- `int`: Status code (0 = success, non-zero = error code) + +**Example:** +```python +result = store.remove("my_key") +if result == 0: + print("Successfully removed") +``` + +--- + +#### remove_by_regex() +Remove objects from the storage system whose keys match a regular expression. + +```python +def remove_by_regex(self, regex: str) -> int +``` + +**Parameters:** +- `regex` (str): The regular expression to match against object keys. + +**Returns:** +- `int`: The number of objects removed, or a negative value on error. + +**Example:** +```python +# Remove all keys starting with "user_session_" +count = store.remove_by_regex("^user_session_.*") +if count >= 0: + print(f"Removed {count} objects") +``` + +--- + +#### remove_all() +Remove all objects from the storage system. + +```python +def remove_all(self) -> int +``` + +**Returns:** +- `int`: Number of objects removed, or -1 on error + +**Example:** +```python +count = store.remove_all() +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. + +```python +def is_exist(self, key: str) -> int +``` + +**Parameters:** +- `key` (str): Object identifier to check + +**Returns:** +- `int`: + - `1`: Object exists + - `0`: Object doesn't exist + - `-1`: Error occurred + +**Example:** +```python +exists = store.is_exist("my_key") +if exists == 1: + print("Object exists") +elif exists == 0: + print("Object not found") +else: + print("Error checking existence") +``` + +--- + +#### batch_is_exist() +Check existence of multiple objects in a single batch operation. + +```python +def batch_is_exist(self, keys: List[str]) -> List[int] +``` + +**Parameters:** +- `keys` (List[str]): List of object identifiers to check + +**Returns:** +- `List[int]`: List of existence results (1=exists, 0=not exists, -1=error) + +**Example:** +```python +keys = ["key1", "key2", "key3"] +results = store.batch_is_exist(keys) +for key, exists in zip(keys, results): + status = "exists" if exists == 1 else "not found" if exists == 0 else "error" + print(f"{key}: {status}") +``` + +--- + +#### get_size() +Get the size of a stored object in bytes. + +```python +def get_size(self, key: str) -> int +``` + +**Parameters:** +- `key` (str): Object identifier + +**Returns:** +- `int`: Size in bytes, or negative value on error + +**Example:** +```python +size = store.get_size("my_key") +if size >= 0: + print(f"Object size: {size} bytes") +else: + print("Error getting size or object not found") +``` + +--- + +#### get_buffer() +Get object data as a buffer that implements Python's buffer protocol. + +```python +def get_buffer(self, key: str) -> BufferHandle +``` + +**Parameters:** +- `key` (str): Object identifier + +**Returns:** +- `BufferHandle`: Buffer object or None if not found + +**Example:** +```python +buffer = store.get_buffer("large_object") +if buffer: + print(f"Buffer size: {buffer.size()} bytes") + # Use with numpy without copying + import numpy as np + arr = np.array(buffer, copy=False) +``` + +--- + +#### put_parts() +Store data from multiple buffer parts as a single object. + +```python +def put_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!" +part3 = b" From Mooncake" +result = store.put_parts("greeting", part1, part2, part3) +``` + +--- +#### batch_get_buffer() +Get multiple objects as buffers that implement Python's buffer protocol. + +```python +def batch_get_buffer(self, keys: List[str]) -> List[BufferHandle] +``` + +**Parameters:** +- `keys` (List[str]): List of object identifiers to retrieve + +**Returns:** +- `List[BufferHandle]`: List of buffer objects, with None for keys not found + +**Note:** This function is supported for dummy clients through the real +client-owned shared-memory staging buffer. + +**Example:** +```python +buffers = store.batch_get_buffer(["key1", "key2", "key3"]) +for i, buffer in enumerate(buffers): + if buffer: + print(f"Buffer {i} size: {buffer.size()} bytes") +``` + +--- +#### alloc_from_mem_pool() +Allocate memory from the memory pool. + +```python +def alloc_from_mem_pool(self, size: int) -> int +``` + +**Parameters:** +- `size` (int): Size of memory to allocate in bytes + +**Returns:** +- `int`: Memory address as integer, or 0 on failure + +--- +#### init_all() +Initialize all resources with specified protocol and device. + +```python +def init_all(self, protocol: str, device_name: str, mount_segment_size: int = 16777216) -> int +``` + +**Parameters:** +- `protocol` (str): Network protocol - "tcp" or "rdma" +- `device_name` (str): Device name for the protocol +- `mount_segment_size` (int): Memory segment size in bytes for mounting (default: 16MB = 16777216) + +**Returns:** +- `int`: Status code (0 = success, non-zero = error code) + +--- + +#### mount_segment() +Mount a local file or shared-memory region as one or more Mooncake store +segments. + +```python +def mount_segment( + self, + path: str, + size: int, + offset: int = 0, + protocol: str = "tcp", + location: str = "", +) -> dict +``` + +**Parameters:** +- `path` (str): File or shared-memory path to map and mount. +- `size` (int): Number of bytes to mount. +- `offset` (int, optional): File offset in bytes. Defaults to `0`. +- `protocol` (str, optional): Transfer protocol. Defaults to `"tcp"`. +- `location` (str, optional): Device or locality hint. Defaults to an empty + string. + +**Returns:** +- `dict`: A result dictionary with: + - `ret` (int): Status code (0 = success, non-zero = error code) + - `segment_ids` (List[str]): Segment ids created by the mount operation + +**Example:** +```python +result = store.mount_segment( + "/dev/shm/mooncake_segment", + 16 * 1024 * 1024, + offset=0, + protocol="tcp", + location="", +) + +if result["ret"] == 0: + segment_ids = list(result["segment_ids"]) + print("Mounted segments:", segment_ids) +else: + print("Mount failed:", result["ret"]) +``` + +The corresponding HTTP endpoints are `/api/mount_shm` and +`/api/unmount_shm`, but the HTTP API is intentionally narrower: it accepts a +named shared memory object name instead of an arbitrary path. + +--- + +#### unmount_segment() +Unmount one or more file or shared-memory segments by segment id. + +```python +def unmount_segment( + self, + segment_ids: List[str], + grace_period_seconds: int = 0, +) -> int +``` + +**Parameters:** +- `segment_ids` (List[str]): Segment ids returned by `mount_segment()`. +- `grace_period_seconds` (int, optional): Grace period before the segment is + fully unmounted. Defaults to `0`, which keeps the existing immediate unmount + behavior. During a positive grace period, the segment remains readable but no + longer accepts new allocations. + +**Returns:** +- `int`: Status code (0 = success, non-zero = error code) + +**Example:** +```python +ret = store.unmount_segment(segment_ids, grace_period_seconds=30) +if ret != 0: + print("Unmount failed:", ret) +``` + +--- + +#### allocate_and_mount_segment() +Allocate memory inside the store process and mount it as one or more Mooncake +store segments. + +```python +def allocate_and_mount_segment( + self, + size: int, + protocol: str = "tcp", + location: str = "", +) -> dict +``` + +**Parameters:** +- `size` (int): Number of bytes requested. The allocated size may be rounded up + for alignment. +- `protocol` (str, optional): Transfer protocol. Defaults to `"tcp"`. +- `location` (str, optional): Device or locality hint. Defaults to an empty + string. + +**Returns:** +- `dict`: A result dictionary with: + - `ret` (int): Status code (0 = success, non-zero = error code) + - `segment_ids` (List[str]): Segment ids created by the mount operation + - `allocated_size` (int): Actual allocated size in bytes + +**Example:** +```python +result = store.allocate_and_mount_segment( + 16 * 1024 * 1024, + protocol="tcp", + location="", +) + +if result["ret"] == 0: + segment_ids = list(result["segment_ids"]) + allocated_size = result["allocated_size"] +``` + +The corresponding HTTP endpoints are `/api/mount` and `/api/unmount`. + +--- + +#### unmount_and_free_segment() +Unmount one or more internally allocated segments by segment id and free their +local memory. + +```python +def unmount_and_free_segment( + self, + segment_ids: List[str], + grace_period_seconds: int = 0, +) -> int +``` + +**Parameters:** +- `segment_ids` (List[str]): Segment ids returned by + `allocate_and_mount_segment()`. +- `grace_period_seconds` (int, optional): Grace period before the segment is + fully unmounted and its local allocated memory is released. Defaults to `0`, + which keeps the existing immediate unmount-and-free behavior. During a + positive grace period, the segment remains readable but no longer accepts new + allocations. + +**Returns:** +- `int`: Status code (0 = success, non-zero = error code) + +**Example:** +```python +ret = store.unmount_and_free_segment(segment_ids, grace_period_seconds=30) +if ret != 0: + print("Unmount and free failed:", ret) +``` + +--- + +#### get_hostname() +Get the hostname of the current store instance. + +```python +def get_hostname(self) -> str +``` + +**Returns:** +- `str`: Hostname and port of this store instance + +**Example:** +```python +hostname = store.get_hostname() +print(f"Store running on: {hostname}") +``` + +--- + +#### get_replica_desc() +Get descriptors of replicas for a key. + +```python +def get_replica_desc(self, key: str) -> List[Replica::Descriptor] +``` + +**Parameters:** +- `key` (str): mooncake store key + +**Returns:** +- `List[Replica::Descriptor]`: List of replica descriptors + +**Example:** +```python +descriptors = store.get_replica_desc("mooncake_key") +for desc in descriptors: + print("Status:", desc.status) + if desc.is_memory_replica(): + mem_desc = desc.get_memory_descriptor() + print("Memory buffer desc:", mem_desc.buffer_descriptor) + elif desc.is_disk_replica(): + disk_desc = desc.get_disk_descriptor() + print("Disk path:", disk_desc.file_path, "Size:", disk_desc.object_size) +``` + +--- + +#### batch_get_replica_desc() +Get descriptors of replicas for a tuple of keys. + +```python +def batch_get_replica_desc(self, keys: List[str]) -> Dict[str, List[Replica::Descriptor]] +``` + +**Parameters:** +- `keys` (List[str]): List of mooncake store keys + +**Returns:** +- `Dict[str, List[Replica::Descriptor]]`: Dictionary mapping keys to their list of replica descriptors + +**Example:** +```python +descriptors_map = store.batch_get_replica_desc(["key1", "key2"]) +for key, desc_list in descriptors_map.items(): + print(f"Replicas for key: {key}") + for desc in desc_list: + if desc.is_memory_replica(): + mem_desc = desc.get_memory_descriptor() + print("Memory buffer desc:", mem_desc.buffer_descriptor) + elif desc.is_disk_replica(): + disk_desc = desc.get_disk_descriptor() + print("Disk path:", disk_desc.file_path, "Size:", disk_desc.object_size) +``` + +--- + +#### create_copy_task() + +Creates an asynchronous copy task to replicate an object to target segments. + +```python +def create_copy_task(self, key: str, targets: List[str]) -> Tuple[UUID, int] +``` + +**Parameters:** +- `key` (str): Object key to copy +- `targets` (List[str]): List of target segment names where replicas should be created + +**Returns:** +- `Tuple[UUID, int]`: (task UUID, error code) + - If successful: (task UUID, 0) + - If failed: (UUID{0, 0}, error code) + +**Task lifecycle and failure behavior:** +- New tasks start in `TaskStatus.PENDING`, move to `TaskStatus.PROCESSING` after a client picks them up, and finish as `TaskStatus.SUCCESS` or `TaskStatus.FAILED`. +- The task payload is executed by a storage client in the background; the client reports the final status back to the master automatically. +- Only allocation-pressure failures (`NO_AVAILABLE_HANDLE`) are retried automatically by the client, up to the master-side `max_retry_attempts` setting. +- Submission can fail immediately with errors such as `TASK_PENDING_LIMIT_EXCEEDED` when the master-side pending queue is full. + +**Example:** +```python +# Create an asynchronous copy task +task_id, error_code = store.create_copy_task("my_key", ["segment1", "segment2"]) +if error_code == 0: + print(f"Copy task created with ID: {task_id}") + # Query task status later + response, status = store.query_task(task_id) + if status == 0: + print(f"Task status: {response.status}") +else: + print(f"Failed to create copy task: {error_code}") +``` + +--- + +#### create_move_task() + +Creates an asynchronous move task to move an object from source segment to target segment. + +```python +def create_move_task(self, key: str, source: str, target: str) -> Tuple[UUID, int] +``` + +**Parameters:** +- `key` (str): Object key to move +- `source` (str): Source segment name where the replica currently exists +- `target` (str): Target segment name where the replica should be moved to + +**Returns:** +- `Tuple[UUID, int]`: (task UUID, error code) + - If successful: (task UUID, 0) + - If failed: (UUID{0, 0}, error code) + +**Task lifecycle and failure behavior:** +- Move tasks use the same state machine as copy tasks: `PENDING -> PROCESSING -> SUCCESS/FAILED`. +- **Submission-time failures**: If the object or source replica is already missing when `create_move_task` is called, the call returns an error code directly (e.g., `OBJECT_NOT_FOUND` or `INVALID_PARAMS`) and no task is created. +- **Execution-time failures**: If the object or source replica disappears after the task is successfully submitted, the task transitions to `FAILED` state. Only `NO_AVAILABLE_HANDLE` execution failures are retried automatically. +- Timeout and retry behavior is controlled on the master side rather than by the Python client API. + +**Example:** +```python +# Create an asynchronous move task +task_id, error_code = store.create_move_task("my_key", "old_segment", "new_segment") +if error_code == 0: + print(f"Move task created with ID: {task_id}") + # Query task status later + response, status = store.query_task(task_id) + if status == 0: + print(f"Task status: {response.status}") +else: + print(f"Failed to create move task: {error_code}") +``` + +--- + +#### query_task() + +Queries the status of an asynchronous task (copy or move). + +```python +def query_task(self, task_id: UUID) -> Tuple[QueryTaskResponse | None, int] +``` + +**Parameters:** +- `task_id` (UUID): UUID of the task to query + +**Returns:** +- `Tuple[QueryTaskResponse | None, int]`: (QueryTaskResponse if success, error code) + - If successful: (QueryTaskResponse, 0) + - If failed: (None, error code) + +`QueryTaskResponse` includes: +- `id`: task UUID +- `type`: `TaskType.REPLICA_COPY` or `TaskType.REPLICA_MOVE` +- `status`: `TaskStatus.PENDING`, `TaskStatus.PROCESSING`, `TaskStatus.SUCCESS`, or `TaskStatus.FAILED` +- `created_at_ms_epoch`: creation timestamp in milliseconds +- `last_updated_at_ms_epoch`: last state-change timestamp in milliseconds +- `assigned_client`: UUID of the client currently assigned to the task +- `message`: completion or failure message + +Typical query-time failures include: +- `TASK_NOT_FOUND`: the task ID does not exist or the finished task has already been pruned from the master's in-memory history + +**Master-side task manager settings affecting task APIs:** +- `--max_total_finished_tasks`: number of completed tasks retained for later `query_task` calls +- `--max_total_pending_tasks`: maximum queued tasks before submissions fail with `TASK_PENDING_LIMIT_EXCEEDED` +- `--max_total_processing_tasks`: cap on concurrently processing tasks +- `--pending_task_timeout_sec`: how long a task may stay in `PENDING` before being failed by the master (`0` disables this timeout) +- `--processing_task_timeout_sec`: how long a task may stay in `PROCESSING` before being failed by the master (`0` disables this timeout) +- `--max_retry_attempts`: retry budget used only for `NO_AVAILABLE_HANDLE` execution failures + +**Example:** +```python +from mooncake.store import MooncakeDistributedStore, TaskStatus +import time + +# Initialize store +store = MooncakeDistributedStore() +store.setup("localhost", "http://localhost:8080/metadata", + 512*1024*1024, 128*1024*1024, "tcp", "", "localhost:50051") + +# Submit multiple copy tasks +tasks = [] +for key in ["key1", "key2", "key3"]: + task_id, error = store.create_copy_task(key, ["segment1", "segment2"]) + if error == 0: + tasks.append(task_id) + print(f"Created copy task {task_id} for {key}") + +# Monitor task progress +while tasks: + completed = [] + for task_id in tasks: + response, status = store.query_task(task_id) + if status == 0 and response: + if response.status == TaskStatus.SUCCESS: + print(f"Task {task_id} succeeded") + completed.append(task_id) + elif response.status == TaskStatus.FAILED: + print(f"Task {task_id} failed: {response.message}") + completed.append(task_id) + + # Remove completed tasks + tasks = [t for t in tasks if t not in completed] + + if tasks: + time.sleep(1) # Wait before next check + +print("All tasks completed") +store.close() +``` + +--- + +#### close() +Clean up all resources and terminate connections. + +```python +def close(self) -> int +``` + +**Returns:** +- `int`: Status code (0 = success, non-zero = error code) + +**Example:** +```python +store.close() +``` + +--- +#### put_from_with_metadata() +Store data directly from a registered buffer with metadata (zero-copy). + +```python +def put_from_with_metadata(self, key: str, buffer_ptr: int, metadata_buffer_ptr: int, size: int, metadata_size: int, config: ReplicateConfig = None) -> int +``` + +**Parameters:** +- `key` (str): Object identifier +- `buffer_ptr` (int): Memory address of the main data buffer (from ctypes.data or similar) +- `metadata_buffer_ptr` (int): Memory address of the metadata buffer +- `size` (int): Number of bytes for the main data +- `metadata_size` (int): Number of bytes for the metadata +- `config` (ReplicateConfig, optional): Replication configuration + +**Returns:** +- `int`: Status code (0 = success, non-zero = error code) + +**Note:** This function is supported for dummy clients through the real +client-owned shared-memory staging buffer. + +**Example:** +```python +import numpy as np + +# Create data and metadata +data = np.random.randn(1000).astype(np.float32) +metadata = np.array([42, 100], dtype=np.int32) # example metadata + +# Register buffers +data_ptr = data.ctypes.data +metadata_ptr = metadata.ctypes.data +store.register_buffer(data_ptr, data.nbytes) +store.register_buffer(metadata_ptr, metadata.nbytes) + +# Store with metadata +result = store.put_from_with_metadata("data_with_metadata", data_ptr, metadata_ptr, + data.nbytes, metadata.nbytes) +if result == 0: + print("Data with metadata stored successfully") + +# Cleanup +store.unregister_buffer(data_ptr) +store.unregister_buffer(metadata_ptr) +``` + +--- +#### pub_tensor() +Publish a PyTorch tensor with configurable replication settings. + +```python +def pub_tensor(self, key: str, tensor: torch.Tensor, config: ReplicateConfig = None) -> int +``` + +**Parameters:** +- `key` (str): Unique object identifier +- `tensor` (torch.Tensor): PyTorch tensor to store +- `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. + +**Example:** +```python +import torch +from mooncake.store import ReplicateConfig, SoftPinAction + +# Create a tensor +tensor = torch.randn(100, 100) + +# Create replication config +config = ReplicateConfig() +config.replica_num = 3 +config.soft_pin_action = SoftPinAction.ENABLE + +# Publish tensor with replication settings +result = store.pub_tensor("my_tensor", tensor, config) +if result == 0: + print("Tensor published successfully") +``` + +--- + +### PyTorch Tensor Operations (Tensor Parallelism) + +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. + +⚠️ **Note**: These methods require `torch` to be installed and available in the environment. + +#### put_tensor_with_tp() + +Put a PyTorch tensor into the store, optionally splitting it into shards for tensor parallelism. +The tensor is chunked immediately and stored as separate keys (e.g., `key_tp_0`, `key_tp_1`...). + +```python +def put_tensor_with_tp(self, key: str, tensor: torch.Tensor, tp_rank: int = 0, tp_size: int = 1, split_dim: int = 0) -> int +``` + +**Parameters:** + + - `key` (str): Base identifier for the tensor. + - `tensor` (torch.Tensor): The PyTorch tensor to store. + - `tp_rank` (int): Current tensor parallel rank (default: 0). *Note: The method splits and stores all chunks for all ranks regardless of this value.* + - `tp_size` (int): Total tensor parallel size (default: 1). If \> 1, the tensor is split into `tp_size` chunks. + - `split_dim` (int): The dimension to split the tensor along (default: 0). + +**Returns:** + + - `int`: Status code (0 = success, non-zero = error code). + +#### pub_tensor_with_tp() + +Publish a PyTorch tensor into the store with configurable replication settings, optionally splitting it into shards for tensor parallelism. +The tensor is chunked immediately and stored as separate keys (e.g., `key_tp_0`, `key_tp_1`...). + +```python +def pub_tensor_with_tp(self, key: str, tensor: torch.Tensor, config: ReplicateConfig, tp_rank: int = 0, tp_size: int = 1, split_dim: int = 0) -> int +``` + +**Parameters:** + + - `key` (str): Base identifier for the tensor. + - `tensor` (torch.Tensor): The PyTorch tensor to store. + - `config` (ReplicateConfig): Optional replication configuration. + - `tp_rank` (int): Current tensor parallel rank (default: 0). *Note: The method splits and stores all chunks for all ranks regardless of this value.* + - `tp_size` (int): Total tensor parallel size (default: 1). If \> 1, the tensor is split into `tp_size` chunks. + - `split_dim` (int): The dimension to split the tensor along (default: 0). + +**Returns:** + + - `int`: Status code (0 = success, non-zero = error code). + +#### get_tensor_with_tp() + +Get a PyTorch tensor from the store, specifically retrieving the shard corresponding to the given Tensor Parallel rank. + +```python +def get_tensor_with_tp(self, key: str, tp_rank: int = 0, tp_size: int = 1, split_dim: int = 0) -> torch.Tensor +``` + +**Parameters:** + + - `key` (str): Base identifier of the tensor. + - `tp_rank` (int): The tensor parallel rank to retrieve (default: 0). Fetches key `key_tp_{rank}` if `tp_size > 1`. + - `tp_size` (int): Total tensor parallel size (default: 1). + - `split_dim` (int): The dimension used during splitting (default: 0). + +**Returns:** + + - `torch.Tensor`: The retrieved tensor (or shard). Returns `None` if not found. + +#### batch_put_tensor_with_tp() + +Put a batch of PyTorch tensors into the store, splitting each into shards for tensor parallelism. + +```python +def batch_put_tensor_with_tp(self, base_keys: List[str], tensors_list: List[torch.Tensor], tp_rank: int = 0, tp_size: int = 1, split_dim: int = 0) -> List[int] +``` + +**Parameters:** + + - `base_keys` (List[str]): List of base identifiers. + - `tensors_list` (List[torch.Tensor]): List of tensors to store. + - `tp_rank` (int): Current rank (default: 0). + - `tp_size` (int): Total TP size (default: 1). + - `split_dim` (int): Split dimension (default: 0). + +**Returns:** + + - `List[int]`: List of status codes for each tensor operation. + +#### batch_pub_tensor_with_tp() + +Publish a batch of PyTorch tensors into the store with configurable replication settings, splitting each into shards for tensor parallelism. + +```python +def batch_pub_tensor_with_tp(self, base_keys: List[str], tensors_list: List[torch.Tensor], config: ReplicateConfig, tp_rank: int = 0, tp_size: int = 1, split_dim: int = 0) -> List[int] +``` + +**Parameters:** + + - `base_keys` (List[str]): List of base identifiers. + - `tensors_list` (List[torch.Tensor]): List of tensors to store. + - `config` (ReplicateConfig): Optional replication configuration. + - `tp_rank` (int): Current rank (default: 0). + - `tp_size` (int): Total tp size (default: 1). + - `split_dim` (int): Split dimension (default: 0). + +**Returns:** + + - `List[int]`: List of status codes for each tensor operation. + +#### batch_get_tensor_with_tp() + +Get a batch of PyTorch tensor shards from the store for a given Tensor Parallel rank. + +```python +def batch_get_tensor_with_tp(self, base_keys: List[str], tp_rank: int = 0, tp_size: int = 1) -> List[torch.Tensor] +``` + +**Parameters:** + + - `base_keys` (List[str]): List of base identifiers. + - `tp_rank` (int): The tensor parallel rank to retrieve (default: 0). + - `tp_size` (int): Total tensor parallel size (default: 1). + +**Returns:** + + - `List[torch.Tensor]`: List of retrieved tensors (or shards). Contains `None` for missing keys. + +--- +#### put_tensor() + +Put a PyTorch tensor into the store. + +```python +def put_tensor(self, key: str, tensor: torch.Tensor) -> int +``` + +**Parameters:** +- `key` (str): Object identifier +- `tensor` (torch.Tensor): The PyTorch tensor to store + +**Returns:** +- `int`: Status code (0 = success, non-zero = error code) + +**Note:** This function requires `torch` to be installed and available in the environment. + +**Example:** +```python +import torch +from mooncake.store import MooncakeDistributedStore + +store = MooncakeDistributedStore() +store.setup("localhost", "http://localhost:8080/metadata", 512*1024*1024, 128*1024*1024, "tcp", "", "localhost:50051") + +# Store a tensor +tensor = torch.randn(100, 100) +result = store.put_tensor("my_tensor", tensor) +if result == 0: + print("Tensor stored successfully") +``` + +--- + +#### get_tensor() + +Get a PyTorch tensor from the store. + +```python +def get_tensor(self, key: str) -> torch.Tensor +``` + +**Parameters:** +- `key` (str): Object identifier to retrieve + +**Returns:** +- `torch.Tensor`: The retrieved tensor. Returns `None` if not found. + +**Note:** This function requires `torch` to be installed and available in the environment. + +**Example:** +```python +import torch +from mooncake.store import MooncakeDistributedStore + +store = MooncakeDistributedStore() +store.setup("localhost", "http://localhost:8080/metadata", 512*1024*1024, 128*1024*1024, "tcp", "", "localhost:50051") + +# Store a tensor +tensor = torch.randn(100, 100) +store.put_tensor("my_tensor", tensor) + +# Retrieve the tensor +retrieved_tensor = store.get_tensor("my_tensor") +if retrieved_tensor is not None: + print(f"Retrieved tensor with shape: {retrieved_tensor.shape}") +``` + +--- + +#### batch_get_tensor() + +Get a batch of PyTorch tensors from the store. + +```python +def batch_get_tensor(self, keys: List[str]) -> List[torch.Tensor] +``` + +**Parameters:** +- `keys` (List[str]): List of object identifiers to retrieve + +**Returns:** +- `List[torch.Tensor]`: List of retrieved tensors. Contains `None` for missing keys. + +**Note:** This function requires `torch` to be installed and available in the environment. + +**Example:** +```python +import torch +from mooncake.store import MooncakeDistributedStore + +store = MooncakeDistributedStore() +store.setup("localhost", "http://localhost:8080/metadata", 512*1024*1024, 128*1024*1024, "tcp", "", "localhost:50051") + +# Store tensors +tensor1 = torch.randn(100, 100) +tensor2 = torch.randn(50, 50) +store.put_tensor("tensor1", tensor1) +store.put_tensor("tensor2", tensor2) + +# Retrieve multiple tensors +tensors = store.batch_get_tensor(["tensor1", "tensor2", "nonexistent"]) +for i, tensor in enumerate(tensors): + if tensor is not None: + print(f"Tensor {i} shape: {tensor.shape}") + else: + print(f"Tensor {i} not found") +``` + +--- + +#### batch_put_tensor() + +Put a batch of PyTorch tensors into the store. + +```python +def batch_put_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 store + +**Returns:** +- `List[int]`: List of status codes for each tensor operation. + +**Note:** This function requires `torch` to be installed and available in the environment. + +**Example:** +```python +import torch +from mooncake.store import MooncakeDistributedStore + +store = MooncakeDistributedStore() +store.setup("localhost", "http://localhost:8080/metadata", 512*1024*1024, 128*1024*1024, "tcp", "", "localhost:50051") + +# Create tensors +tensors = [torch.randn(100, 100), torch.randn(50, 50), torch.randn(25, 25)] +keys = ["tensor1", "tensor2", "tensor3"] + +# Store multiple tensors +results = store.batch_put_tensor(keys, tensors) +for i, result in enumerate(results): + if result == 0: + print(f"Tensor {i} stored successfully") + else: + print(f"Tensor {i} failed to store with code: {result}") +``` + +#### batch_pub_tensor() + +Pub a batch of PyTorch tensors into the store with configurable replication settings. + +```python +def batch_pub_tensor(self, keys: List[str], tensors_list: List[torch.Tensor], config: ReplicateConfig) -> List[int] +``` + +**Parameters:** + - `keys` (List[str]): List of object identifiers + - `tensors_list` (List[torch.Tensor]): List of tensors to store + - `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. + +--- + +#### 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 +`[TensorObjectHeader+layout metadata][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 supported for dummy clients through the real +client-owned shared-memory staging buffer. + +#### batch_upsert_tensor_from() + +Upsert multiple tensors directly from pre-allocated buffers. Each buffer must +use layout `[TensorObjectHeader+layout metadata][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. It is supported for dummy clients through the real client-owned +shared-memory staging buffer. + +#### 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. It is supported for dummy clients through the real client-owned +shared-memory staging buffer. + +**Example:** +```python +import torch +from mooncake.store import ReplicateConfig, SoftPinAction + +tensor = torch.randn(100, 100) + +config = ReplicateConfig() +config.replica_num = 2 +config.soft_pin_action = SoftPinAction.ENABLE + +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. It is supported for dummy clients through the real client-owned +shared-memory staging buffer. + + +--- + +### 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. + +⚠️ **Note**: These methods require `torch` to be installed and available in the environment. + +#### get_tensor_into() + +Get a PyTorch tensor from the store directly into a pre-allocated buffer. + +```python +def get_tensor_into(self, key: str, buffer_ptr: int, size: int) -> torch.Tensor +``` + +**Parameters:** + + - `key` (str): Base identifier of the tensor. + - `buffer_ptr` (int): The buffer pointer pre-allocated for tensor, and the buffer should be registered. + - `size` (int): The size of buffer. + +**Returns:** + + - `torch.Tensor`: The retrieved tensor (or shard). Returns `None` if not found. + +#### batch_get_tensor() + +Get a batch of PyTorch tensor from the store directly into a pre-allocated buffer. + +```python +def batch_get_tensor_into(self, base_keys: List[str], buffer_ptrs: List[int], sizes: List[int]) -> List[torch.Tensor] +``` + +**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. + +**Returns:** + + - `List[torch.Tensor]`: List of retrieved tensors (or shards). Contains `None` for missing keys. + +#### get_tensor_with_tp_into() + +Get a PyTorch tensor from the store, specifically retrieving the shard corresponding to the given Tensor Parallel rank, directly into the pre-allocated buffer. + +```python +def get_tensor_with_tp_into(self, key: str, buffer_ptr: int, size: int, tp_rank: int = 0, tp_size: int = 1, split_dim: int = 0) -> torch.Tensor +``` + +**Parameters:** + + - `key` (str): Base identifier of the tensor. + - `buffer_ptr` (int): The buffer pointer pre-allocated for tensor, and the buffer should be registered. + - `size` (int): The size of buffer. + - `tp_rank` (int): The tensor parallel rank to retrieve (default: 0). Fetches key `key_tp_{rank}` if `tp_size > 1`. + - `tp_size` (int): Total tensor parallel size (default: 1). + - `split_dim` (int): The dimension used during splitting (default: 0). + +**Returns:** + + - `torch.Tensor`: The retrieved tensor (or shard). Returns `None` if not found. + +#### batch_get_tensor_with_tp_into() + +Get a batch of PyTorch tensor shards from the store for a given Tensor Parallel rank, directly into the pre-allocated buffer. + +```python +def batch_get_tensor_with_tp_into(self, base_keys: List[str], buffer_ptrs: List[int], sizes: List[int], tp_rank: int = 0, tp_size: int = 1) -> List[torch.Tensor] +``` + +**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. + - `tp_rank` (int): The tensor parallel rank to retrieve (default: 0). + - `tp_size` (int): Total tensor parallel size (default: 1). + +**Returns:** + + - `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_`. + +```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 + +#### batch_put_from() +Store multiple objects from pre-registered buffers (zero-copy). + +```python +def batch_put_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 memory addresses +- `sizes` (List[int]): List of buffer sizes +- `config` (ReplicateConfig, optional): Replication configuration + +**Returns:** +- `List[int]`: List of status codes for each operation (0 = success, negative = error) + +--- + +#### batch_get_into() +Retrieve multiple objects into pre-registered buffers (zero-copy). + +```python +def batch_get_into(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 memory addresses +- `sizes` (List[int]): List of buffer sizes + +**Returns:** +- `List[int]`: List of bytes read for each operation (positive = success, negative = error) + +⚠️ **Store-managed Buffer Required**: All buffers must resolve to Store-managed registered memory before batch zero-copy operations. + +**Example:** + +
+Click to expand: Batch zero-copy retrieval example + +```python +# Prepare buffers +keys = ["tensor1", "tensor2", "tensor3"] +buffer_size = 1024 * 1024 # 1MB each +buffers = [] +buffer_ptrs = [] + +for i in range(len(keys)): + buffer = np.empty(buffer_size, dtype=np.uint8) + buffers.append(buffer) + buffer_ptrs.append(buffer.ctypes.data) + store.register_buffer(buffer.ctypes.data, buffer_size) + +# Batch retrieve +sizes = [buffer_size] * len(keys) +results = store.batch_get_into(keys, buffer_ptrs, sizes) + +# Check results +for key, result in zip(keys, results): + if result > 0: + print(f"Retrieved {key}: {result} bytes") + else: + print(f"Failed to retrieve {key}: error {result}") + +# Cleanup +for ptr in buffer_ptrs: + store.unregister_buffer(ptr) +``` + +
+ +--- + +#### batch_put_from_multi_buffers() +Store multiple objects from multiple pre-registered buffers (zero-copy). + +```python +def batch_put_from_multi_buffers(self, keys: List[str], all_buffer_ptrs: List[List[int]], all_sizes: List[List[int]], + config: ReplicateConfig = None) -> List[int] +``` + +**Parameters:** +- `keys` (List[str]): List of object identifiers +- `all_buffer_ptrs` (List[int]): all List of memory addresses +- `sizes` (List[int]): all List of buffer sizes +- `config` (ReplicateConfig, optional): Replication configuration + +**Returns:** +- `List[int]`: List of status codes for each operation (0 = success, negative = error) + +--- + +#### batch_get_into_multi_buffers() + +Retrieve multiple objects into multiple pre-registered buffers (zero-copy). + +```python +def batch_get_into_multi_buffers(self, keys: List[str], all_buffer_ptrs: List[List[int], all_sizes: List[List[int]) -> +List[int] +``` + +**Parameters:** +- `keys` (List[str]): List of object identifiers +- `all_buffer_ptrs` (List[int]): List of memory addresses +- `all_sizes` (List[int]): List of buffer sizes + +**Returns:** +- `List[int]`: List of bytes read for each operation (positive = success, negative = error) + +⚠️ **Store-managed Buffer Required**: All buffers must resolve to Store-managed registered memory before batch zero-copy operations. + +**Example:** + +
+Click to expand: Batch zero-copy put and get for multiple buffers example + +```python +tensor = torch.ones(10, 61, 128*1024, dtype=torch.int8) +data_ptr = tensor.data_ptr() +store.register_buffer(data_ptr, 10*61*128*1024) + +target_tensor = torch.zeros(10, 61, 128*1024, dtype=torch.int8) +target_data_ptr = target_tensor.data_ptr() +store.register_buffer(target_data_ptr, 10*61*128*1024) + +all_local_addrs = [] +all_remote_addrs = [] +all_sizes = [] +keys = [] +for block_i in range(10): + local_addrs = [] + remote_addrs = [] + sizes = [] + for _ in range(61): + local_addrs.append(data_ptr) + remote_addrs.append(target_data_ptr) + sizes.append(128*1024) + data_ptr += 128*1024 + target_data_ptr += 128*1024 + all_local_addrs.append(local_addrs) + all_remote_addrs.append(remote_addrs) + all_sizes.append(sizes) + keys.append(f"kv_{rank}_{block_i}") + +config = ReplicateConfig() +config.prefer_alloc_in_same_node = True +store.batch_put_from_multi_buffers(keys, all_local_addrs, all_sizes, config) +store.batch_get_into_multi_buffers(keys, all_remote_addrs, all_sizes, True) + +store.unregister_buffer(tensor.data_ptr()) +store.unregister_buffer(target_tensor.data_ptr()) +``` +
+ +--- + +### Session-based ranged multi-buffer transfer + +For layerwise KV load/save, resolve Master metadata once per object, then transfer +object-byte ranges across multiple buffers without re-querying Master on every layer. + +Typical flow: + +- Get: `batch_get_session_start` → `batch_get_into_multi_buffer_ranges` (per layer) → `batch_get_session_end` +- Put: `batch_put_session_start` → `batch_put_from_multi_buffer_ranges` (per layer) → `batch_put_session_end` / `batch_put_session_revoke` + +Get sessions cache a filtered `QueryResult` (single complete memory replica + lease). +Range calls only check the cached lease locally (zero Master RPCs). Put sessions +reserve object space via Master `BatchPutStart` and finalize with `BatchPutEnd`. + +Put sessions write MEMORY replicas only. `nof_replica_num > 0` is accepted only for +flexible dual-replica configs (`replica_num == 1` and `nof_replica_num == 1`), where +`batch_put_session_end` finalizes MEMORY and revokes the unused NoF reservation. +Reliable multi-replica NoF configs are rejected at session start. `end` / `revoke` +seal the session (no further range writes) and wait for in-flight range transfers +before talking to Master. + +⚠️ **Store-managed Buffer Required**: All buffers must resolve to Store-managed +registered memory before ranged zero-copy operations. + +#### batch_get_session_start() + +Query replicas once and open a get session for the given keys. + +```python +def batch_get_session_start(self, keys: List[str]) -> List[int] +``` + +**Parameters:** +- `keys` (List[str]): Object identifiers + +**Returns:** +- `List[int]`: Per-key status (0 = success, negative = error) + +#### batch_get_into_multi_buffer_ranges() + +Ranged get into multiple buffers using an active get session (no Master RPC). + +```python +def batch_get_into_multi_buffer_ranges( + self, + keys: List[str], + all_buffer_ptrs: List[List[int]], + all_sizes: List[List[int]], + all_src_offsets: List[List[int]], +) -> List[int] +``` + +**Parameters:** +- `keys` (List[str]): Object identifiers (must have an active get session) +- `all_buffer_ptrs` (List[List[int]]): Per-key list of destination buffer addresses +- `all_sizes` (List[List[int]]): Per-key list of transfer sizes in bytes +- `all_src_offsets` (List[List[int]]): Per-key list of object-byte source offsets + +**Returns:** +- `List[int]`: Bytes transferred per key (positive = success, negative = error) + +#### batch_get_session_end() + +Drop cached get-session metadata for the given keys. + +```python +def batch_get_session_end(self, keys: List[str]) -> int +``` + +**Parameters:** +- `keys` (List[str]): Object identifiers + +**Returns:** +- `int`: 0 on success, negative on error + +#### batch_put_session_start() + +Reserve objects and open a put session without transferring data. + +```python +def batch_put_session_start( + self, + keys: List[str], + sizes: List[int], + config: ReplicateConfig = None, +) -> List[int] +``` + +**Parameters:** +- `keys` (List[str]): Object identifiers +- `sizes` (List[int]): Full object sizes in bytes +- `config` (ReplicateConfig, optional): Replication configuration (applies at start only). + If `group_ids` is set, its length must equal `len(keys)`. When some keys already + have a put session, they are skipped and `group_ids` is filtered to match the + remaining keys. + +**Returns:** +- `List[int]`: Per-key status (0 = success, negative = error) + +#### batch_put_from_multi_buffer_ranges() + +Ranged put from multiple buffers using an active put session (no Master RPC). + +```python +def batch_put_from_multi_buffer_ranges( + self, + keys: List[str], + all_buffer_ptrs: List[List[int]], + all_sizes: List[List[int]], + all_dst_offsets: List[List[int]], +) -> List[int] +``` + +**Parameters:** +- `keys` (List[str]): Object identifiers (must have an active put session) +- `all_buffer_ptrs` (List[List[int]]): Per-key list of source buffer addresses +- `all_sizes` (List[List[int]]): Per-key list of transfer sizes in bytes +- `all_dst_offsets` (List[List[int]]): Per-key list of object-byte destination offsets + +**Returns:** +- `List[int]`: Bytes transferred per key (positive = success, negative = error) + +#### batch_put_session_end() + +Finalize a put session and make objects readable. + +```python +def batch_put_session_end(self, keys: List[str]) -> List[int] +``` + +**Parameters:** +- `keys` (List[str]): Object identifiers + +**Returns:** +- `List[int]`: Per-key status (0 = success, negative = error) + +#### batch_put_session_revoke() + +Abort an incomplete put session and release reserved space. + +```python +def batch_put_session_revoke(self, keys: List[str]) -> List[int] +``` + +**Parameters:** +- `keys` (List[str]): Object identifiers + +**Returns:** +- `List[int]`: Per-key status (0 = success, negative = error) + +**Example:** + +
+Click to expand: Session ranged put/get example + +```python +import numpy as np + +page = 1024 +layers = 4 +keys = ["block0", "block1"] +object_sizes = [page * layers] * len(keys) + +# Prepare one registered buffer per layer for each key +src = [np.full(page, i, dtype=np.uint8) for i in range(layers)] +dst = [np.zeros(page, dtype=np.uint8) for _ in range(layers)] +for buf in src + dst: + store.register_buffer(buf.ctypes.data, buf.nbytes) + +assert all(rc == 0 for rc in store.batch_put_session_start(keys, object_sizes)) +for layer in range(layers): + ptrs = [[src[layer].ctypes.data] for _ in keys] + sizes = [[page] for _ in keys] + offsets = [[layer * page] for _ in keys] + rcs = store.batch_put_from_multi_buffer_ranges(keys, ptrs, sizes, offsets) + assert all(rc == page for rc in rcs) +assert all(rc == 0 for rc in store.batch_put_session_end(keys)) + +assert all(rc == 0 for rc in store.batch_get_session_start(keys)) +for layer in range(layers): + ptrs = [[dst[layer].ctypes.data] for _ in keys] + sizes = [[page] for _ in keys] + offsets = [[layer * page] for _ in keys] + rcs = store.batch_get_into_multi_buffer_ranges(keys, ptrs, sizes, offsets) + assert all(rc == page for rc in rcs) +assert store.batch_get_session_end(keys) == 0 +``` +
+ +## MooncakeHostMemAllocator Class + +The `MooncakeHostMemAllocator` class provides host memory allocation capabilities for Mooncake Store operations. + +### Class Definition + +```python +from mooncake.store import MooncakeHostMemAllocator + +# Create an allocator instance +allocator = MooncakeHostMemAllocator() +``` + +### Methods + +#### alloc() +Allocate memory from the host memory pool. + +```python +def alloc(self, size: int) -> int +``` + +**Parameters:** +- `size` (int): Size of memory to allocate in bytes + +**Returns:** +- `int`: Memory address as integer, or 0 on failure + +**Example:** +```python +allocator = MooncakeHostMemAllocator() +ptr = allocator.alloc(1024 * 1024) # Allocate 1MB +if ptr != 0: + print(f"Allocated memory at address: {ptr}") +``` + +#### free() +Free previously allocated memory. + +```python +def free(self, ptr: int) -> int +``` + +**Parameters:** +- `ptr` (int): Memory address to free + +**Returns:** +- `int`: Status code (0 = success, non-zero = error code) + +**Example:** +```python +result = allocator.free(ptr) +if result == 0: + print("Memory freed successfully") +``` + +--- + +## bind_to_numa_node Function + +The `bind_to_numa_node` function binds the current thread and memory allocation preference to a specified NUMA node. + +### Function Definition + +```python +from mooncake.store import bind_to_numa_node + +# Bind to NUMA node +bind_to_numa_node(node: int) +``` + +**Parameters:** +- `node` (int): NUMA node number to bind to + +**Example:** +```python +from mooncake.store import bind_to_numa_node + +# Bind current thread to NUMA node 0 +bind_to_numa_node(0) +``` + +--- + +## Error Handling + +Most methods return integer status codes: +- `0`: Success +- Negative values: Error codes (for methods that can return data size) + +For methods that return data (`get`, `get_batch`, `get_buffer`, `get_tensor`): +- Return the requested data on success +- Return empty/None on failure or key not found + +📋 **Complete Error Codes Reference**: See [Error Code Explanation](../../troubleshooting/error-code.md) for detailed descriptions of all error codes and their meanings. + +--- + +## Performance Tips + +1. **Use batch operations** when working with multiple objects to reduce network overhead +2. **Use zero-copy APIs** (`put_from`, `get_into`) for large data transfers +3. **Register buffers** once and reuse them for multiple operations +4. **Configure replication** appropriately - more replicas provide better availability but use more storage +5. **Use soft pinning** for frequently accessed objects to keep them in memory +6. **Choose RDMA protocol** when available for maximum performance diff --git a/docs/source/python-api-reference/transfer-engine.md b/docs/source/api-reference/python/transfer-engine.md similarity index 81% rename from docs/source/python-api-reference/transfer-engine.md rename to docs/source/api-reference/python/transfer-engine.md index d23b70f51c..955b7ac257 100644 --- a/docs/source/python-api-reference/transfer-engine.md +++ b/docs/source/api-reference/python/transfer-engine.md @@ -4,7 +4,8 @@ The Transfer Engine Python API provides a high-level interface for efficient data transfer between distributed systems using RDMA (Remote Direct Memory Access) and other transport protocols. It enables fast, low-latency data movement between nodes in a cluster. -For interfaces beyond the Python API (C/C++, Golang, Rust), see [Transfer Engine](../design/transfer-engine/index.md#using-transfer-engine-to-your-projects). +For examples and interfaces beyond the Python API (C/C++, Golang, Rust), see +[Using Transfer Engine in Your Projects](../../design/transfer-engine/index.md#using-transfer-engine-in-your-projects). ## Installation @@ -18,7 +19,8 @@ pip install mooncake-transfer-engine ## Quick Start -See the [Transfer Engine Quick Start](../getting_started/quick-start.md#transfer-engine-quick-start) guide for a complete example of setting up and using the Transfer Engine. +See the [Transfer Engine Python quick start](../../design/transfer-engine/index.md#using-transfer-engine-in-your-projects) +for a complete example of setting up and using the Transfer Engine API directly. ## API Reference @@ -164,10 +166,14 @@ Gets the address of the first buffer in a specified segment. ### Data Transfer Operations +```{note} +The optional `transport_hint` argument pins the request onto a named transport (`"rdma"`, `"tcp"`, ...), overriding policy-driven selection for that one call. **TENT backend only** (`MC_USE_TENT=1`); silently ignored on the classic backend. See [TENT transport selector](../../design/tent/transport-selector.md) for more details. +``` + #### transfer_sync_write() ```python -transfer_sync_write(target_hostname, buffer, peer_buffer_address, length) +transfer_sync_write(target_hostname, buffer, peer_buffer_address, length, transport_hint="") ``` Performs a synchronous write operation to transfer data from local buffer to remote buffer. @@ -177,6 +183,7 @@ Performs a synchronous write operation to transfer data from local buffer to rem - `buffer` (int): The local buffer address - `peer_buffer_address` (int): The remote buffer address - `length` (int): The number of bytes to transfer +- `transport_hint` (str, optional): TENT-only per-request transport pin. See the note above. Default `""` (policy-driven). **Returns:** - `int`: 0 on success, negative value on failure @@ -184,7 +191,7 @@ Performs a synchronous write operation to transfer data from local buffer to rem #### transfer_sync_read() ```python -transfer_sync_read(target_hostname, buffer, peer_buffer_address, length) +transfer_sync_read(target_hostname, buffer, peer_buffer_address, length, transport_hint="") ``` Performs a synchronous read operation to transfer data from remote buffer to local buffer. @@ -194,6 +201,7 @@ Performs a synchronous read operation to transfer data from remote buffer to loc - `buffer` (int): The local buffer address - `peer_buffer_address` (int): The remote buffer address - `length` (int): The number of bytes to transfer +- `transport_hint` (str, optional): TENT-only per-request transport pin. See the note above. Default `""` (policy-driven). **Returns:** - `int`: 0 on success, negative value on failure @@ -201,7 +209,7 @@ Performs a synchronous read operation to transfer data from remote buffer to loc #### transfer_sync() ```python -transfer_sync(target_hostname, buffer, peer_buffer_address, length, opcode, notify=None) +transfer_sync(target_hostname, buffer, peer_buffer_address, length, opcode, notify=None, transport_hint="") ``` Performs a synchronous transfer operation with specified opcode and optional notification. @@ -213,6 +221,7 @@ Performs a synchronous transfer operation with specified opcode and optional not - `length` (int): The number of bytes to transfer - `opcode` (TransferOpcode): The transfer operation type (READ or WRITE) - `notify` (TransferNotify, optional): Notification object to send after transfer completion +- `transport_hint` (str, optional): TENT-only per-request transport pin. See the note above. Default `""` (policy-driven). **Returns:** - `int`: 0 on success, negative value on failure @@ -220,7 +229,7 @@ Performs a synchronous transfer operation with specified opcode and optional not #### transfer_submit_write() ```python -transfer_submit_write(target_hostname, buffer, peer_buffer_address, length) +transfer_submit_write(target_hostname, buffer, peer_buffer_address, length, transport_hint="") ``` Submits an asynchronous write operation and returns immediately. @@ -230,9 +239,10 @@ Submits an asynchronous write operation and returns immediately. - `buffer` (int): The local buffer address - `peer_buffer_address` (int): The remote buffer address - `length` (int): The number of bytes to transfer +- `transport_hint` (str, optional): TENT-only per-request transport pin. See the note above. Default `""` (policy-driven). **Returns:** -- `int`: Batch ID for tracking the operation, or negative value on failure +- `int`: Batch ID for tracking the operation, or 0 on failure #### transfer_check_status() @@ -255,7 +265,7 @@ Checks the status of an asynchronous transfer operation. #### transfer_write_on_cuda() ```python -transfer_write_on_cuda(target_hostname, buffer, peer_buffer_address, length, stream_ptr) +transfer_write_on_cuda(target_hostname, buffer, peer_buffer_address, length, stream_ptr=0, transport_hint="") ``` Performs a write operation to transfer data from local buffer to remote buffer on a given cuda stream. @@ -266,6 +276,7 @@ Performs a write operation to transfer data from local buffer to remote buffer o - `peer_buffer_address` (int): The remote buffer address - `length` (int): The number of bytes to transfer - `stream_ptr` (int): The integer representation of a CUDA stream pointer (`cudaStream_t`). For example, from a PyTorch stream, this can be obtained via `stream.cuda_stream`. +- `transport_hint` (str, optional): TENT-only per-request transport pin. See the note at the top of this section. Default `""` (policy-driven). **Returns:** - `None`: The function returns immediately after successfully scheduling the transfer callback. @@ -279,7 +290,7 @@ Performs a write operation to transfer data from local buffer to remote buffer o #### transfer_read_on_cuda() ```python -transfer_read_on_cuda(target_hostname, buffer, peer_buffer_address, length, stream_ptr) +transfer_read_on_cuda(target_hostname, buffer, peer_buffer_address, length, stream_ptr=0, transport_hint="") ``` Performs a read operation to transfer data from remote buffer to local buffer on a given cuda stream. @@ -290,6 +301,7 @@ Performs a read operation to transfer data from remote buffer to local buffer on - `peer_buffer_address` (int): The remote buffer address - `length` (int): The number of bytes to transfer - `stream_ptr` (int): The integer representation of a CUDA stream pointer (`cudaStream_t`). For example, from a PyTorch stream, this can be obtained via `stream.cuda_stream`. +- `transport_hint` (str, optional): TENT-only per-request transport pin. See the note at the top of this section. Default `""` (policy-driven). **Returns:** - `None`: The function returns immediately after successfully scheduling the transfer callback. @@ -307,7 +319,7 @@ Performs a read operation to transfer data from remote buffer to local buffer on #### batch_transfer_sync_write() ```python -batch_transfer_sync_write(target_hostname, buffers, peer_buffer_addresses, lengths) +batch_transfer_sync_write(target_hostname, buffers, peer_buffer_addresses, lengths, transport_hint="") ``` Performs a batch synchronous write operation to transfer multiple data chunks from local buffers to remote buffers. @@ -317,6 +329,7 @@ Performs a batch synchronous write operation to transfer multiple data chunks fr - `buffers` (List[int]): List of local buffer addresses - `peer_buffer_addresses` (List[int]): List of remote buffer addresses - `lengths` (List[int]): List of byte lengths for each transfer +- `transport_hint` (str, optional): TENT-only per-batch transport pin. Default `""` (policy-driven). **Returns:** - `int`: 0 on success, negative value on failure @@ -324,7 +337,7 @@ Performs a batch synchronous write operation to transfer multiple data chunks fr #### batch_transfer_sync_read() ```python -batch_transfer_sync_read(target_hostname, buffers, peer_buffer_addresses, lengths) +batch_transfer_sync_read(target_hostname, buffers, peer_buffer_addresses, lengths, transport_hint="") ``` Performs a batch synchronous read operation to transfer multiple data chunks from remote buffers to local buffers. @@ -334,6 +347,7 @@ Performs a batch synchronous read operation to transfer multiple data chunks fro - `buffers` (List[int]): List of local buffer addresses - `peer_buffer_addresses` (List[int]): List of remote buffer addresses - `lengths` (List[int]): List of byte lengths for each transfer +- `transport_hint` (str, optional): TENT-only per-batch transport pin. Default `""` (policy-driven). **Returns:** - `int`: 0 on success, negative value on failure @@ -341,7 +355,7 @@ Performs a batch synchronous read operation to transfer multiple data chunks fro #### batch_transfer_sync() ```python -batch_transfer_sync(target_hostname, buffers, peer_buffer_addresses, lengths, opcode, notify=None) +batch_transfer_sync(target_hostname, buffers, peer_buffer_addresses, lengths, opcode, notify=None, transport_hint="") ``` Performs a batch synchronous transfer operation with specified opcode and optional notification. @@ -353,6 +367,7 @@ Performs a batch synchronous transfer operation with specified opcode and option - `lengths` (List[int]): List of byte lengths for each transfer - `opcode` (TransferOpcode): The transfer operation type (READ or WRITE) - `notify` (TransferNotify, optional): Notification object to send after transfer completion +- `transport_hint` (str, optional): TENT-only per-batch transport pin. Default `""` (policy-driven). **Returns:** - `int`: 0 on success, negative value on failure @@ -360,7 +375,7 @@ Performs a batch synchronous transfer operation with specified opcode and option #### batch_transfer_async_write() ```python -batch_transfer_async_write(target_hostname, buffers, peer_buffer_addresses, lengths) +batch_transfer_async_write(target_hostname, buffers, peer_buffer_addresses, lengths, transport_hint="") ``` Submits a batch asynchronous write operation and returns immediately. @@ -370,6 +385,7 @@ Submits a batch asynchronous write operation and returns immediately. - `buffers` (List[int]): List of local buffer addresses - `peer_buffer_addresses` (List[int]): List of remote buffer addresses - `lengths` (List[int]): List of byte lengths for each transfer +- `transport_hint` (str, optional): TENT-only per-batch transport pin. Default `""` (policy-driven). **Returns:** - `int`: Batch ID for tracking the operation, or 0 on failure @@ -377,7 +393,7 @@ Submits a batch asynchronous write operation and returns immediately. #### batch_transfer_async_read() ```python -batch_transfer_async_read(target_hostname, buffers, peer_buffer_addresses, lengths) +batch_transfer_async_read(target_hostname, buffers, peer_buffer_addresses, lengths, transport_hint="") ``` Submits a batch asynchronous read operation and returns immediately. @@ -387,6 +403,7 @@ Submits a batch asynchronous read operation and returns immediately. - `buffers` (List[int]): List of local buffer addresses - `peer_buffer_addresses` (List[int]): List of remote buffer addresses - `lengths` (List[int]): List of byte lengths for each transfer +- `transport_hint` (str, optional): TENT-only per-batch transport pin. Default `""` (policy-driven). **Returns:** - `int`: Batch ID for tracking the operation, or 0 on failure @@ -394,7 +411,7 @@ Submits a batch asynchronous read operation and returns immediately. #### batch_transfer_async() ```python -batch_transfer_async(target_hostname, buffers, peer_buffer_addresses, lengths, opcode) +batch_transfer_async(target_hostname, buffers, peer_buffer_addresses, lengths, opcode, transport_hint="") ``` Submits a batch asynchronous transfer operation with specified opcode and returns immediately. @@ -405,6 +422,7 @@ Submits a batch asynchronous transfer operation with specified opcode and return - `peer_buffer_addresses` (List[int]): List of remote buffer addresses - `lengths` (List[int]): List of byte lengths for each transfer - `opcode` (TransferOpcode): The transfer operation type (READ or WRITE) +- `transport_hint` (str, optional): TENT-only per-batch transport pin. Default `""` (policy-driven). **Returns:** - `int`: Batch ID for tracking the operation, or 0 on failure @@ -426,7 +444,7 @@ Waits for multiple batch asynchronous transfer operations to complete. #### batch_transfer_write_on_cuda() ```python -batch_transfer_write_on_cuda(target_hostname, buffers, peer_buffer_addresses, lengths, stream_ptr) +batch_transfer_write_on_cuda(target_hostname, buffers, peer_buffer_addresses, lengths, stream_ptr=0, transport_hint="") ``` Performs a batch write operation to transfer multiple data chunks from local buffers to remote buffers on a given cuda stream. @@ -437,6 +455,7 @@ Performs a batch write operation to transfer multiple data chunks from local buf - `peer_buffer_addresses` (List[int]): List of remote buffer addresses - `lengths` (List[int]): List of byte lengths for each transfer - `stream_ptr` (int): The integer representation of a CUDA stream pointer (`cudaStream_t`). For example, from a PyTorch stream, this can be obtained via `stream.cuda_stream`. +- `transport_hint` (str, optional): TENT-only per-batch transport pin. Default `""` (policy-driven). **Returns:** - `None`: The function returns immediately after successfully scheduling the transfer callback. @@ -450,7 +469,7 @@ Performs a batch write operation to transfer multiple data chunks from local buf #### batch_transfer_read_on_cuda() ```python -batch_transfer_read_on_cuda(target_hostname, buffers, peer_buffer_addresses, lengths, stream_ptr) +batch_transfer_read_on_cuda(target_hostname, buffers, peer_buffer_addresses, lengths, stream_ptr=0, transport_hint="") ``` Performs a batch read operation to transfer multiple data chunks from remote buffers to local buffers on a given cuda stream. @@ -461,6 +480,7 @@ Performs a batch read operation to transfer multiple data chunks from remote buf - `peer_buffer_addresses` (List[int]): List of remote buffer addresses - `lengths` (List[int]): List of byte lengths for each transfer - `stream_ptr` (int): The integer representation of a CUDA stream pointer (`cudaStream_t`). For example, from a PyTorch stream, this can be obtained via `stream.cuda_stream`. +- `transport_hint` (str, optional): TENT-only per-batch transport pin. Default `""` (policy-driven). **Returns:** - `None`: The function returns immediately after successfully scheduling the transfer callback. @@ -588,6 +608,25 @@ Gets the list of pending transfer notifications received from other nodes. **Returns:** - `List[TransferNotify]`: List of notification objects containing name and message +#### send_probe() + +```python +send_probe(peer_server_name) +``` + +Sends a lightweight JSON-RPC probe to a peer to verify reachability. Used to +test whether a previously-unreachable peer has recovered (e.g. SGLang's +`MooncakeKVManager` uses this to clear entries from its `failed_sessions` +blacklist). + +**Parameters:** +- `peer_server_name` (str): Peer hostname in `host:port` form, as registered + with the metadata server. + +**Returns:** +- `int`: 0 on success, non-zero on failure (peer unknown, unreachable, or + RPC error). + ## Environment Variables The Transfer Engine respects the following environment variables: @@ -596,7 +635,9 @@ The Transfer Engine respects the following environment variables: - `MC_METADATA_SERVER`: Default metadata server address - `MC_LEGACY_RPC_PORT_BINDING`: Enables legacy RPC port binding behavior - `MC_TCP_BIND_ADDRESS`: Specifies the TCP bind address -- `MC_CUSTOM_TOPO_JSON`: Path to custom topology JSON file +- `MC_RDMA_BIND_ADDRESS`: Specifies the RDMA bind address for NIC path construction in dual-NIC environments. When set, RDMA NIC paths use this address while TCP handshake uses the address from `local_hostname`. This is useful when TCP and RDMA traffic use separate network interfaces (e.g., `eth0` for TCP and `rdma-net1` for RDMA). +- `MC_CUSTOM_TOPO_JSON`: Path to custom topology JSON file (classic NIC priority matrix, or under TENT also native `nics`/`mems` JSON). Honored by classic Transfer Engine and by TENT when `MC_USE_TENT=1` (maps to `topology/custom_json_path`). For TENT, prefer inlining `topology/priority_matrix` in `MC_TENT_CONF` when possible; see the {ref}`TENT C++ API `. +- `MC_TE_FILTERS`: Optional comma-separated whitelist of IB device names (e.g. `mlx5_0,mlx5_2`) for legacy Transfer Engine topology discovery. When unset, all available devices are discovered. - `MC_TE_METRIC`: Enables metrics reporting (set to "1", "true", "yes", or "on"). **Note:** Not supported when using Transfer Engine TENT. - `MC_TE_METRIC_INTERVAL_SECONDS`: Sets metrics reporting interval in seconds @@ -605,6 +646,7 @@ The Transfer Engine respects the following environment variables: ### Basic Setup and Data Transfer ```python +import numpy as np from mooncake.engine import TransferEngine import os @@ -631,7 +673,7 @@ data_len = len(data) engine.register_memory(buffer_data, buffer_data_len) # Get Remote Addr from ZMQ or upper-layer inference framework -remote_addr = ?? +remote_addr = REMOTE_ADDR # Transfer data to remote node ret = engine.transfer_sync_write( @@ -796,3 +838,34 @@ The Transfer Engine Python API is thread-safe for most operations. However, it's 2. **Transfer Failures**: Verify target hostname is correct and network connectivity is established 3. **Memory Issues**: Ensure sufficient system memory and proper buffer alignment 4. **Performance Issues**: Check RDMA device configuration and network topology + +## Compile-time Feature Support Attributes + +The `mooncake.engine` module provides boolean attributes that indicate whether specific features were enabled during compilation: + +### Module Attributes + +- `engine.SUPPORT_CUDA`: Whether CUDA support is enabled +- `engine.SUPPORT_EFA`: Whether EFA (Elastic Fabric Adapter) support is enabled +- `engine.SUPPORT_HIP`: Whether HIP (Heterogeneous-compute Interface for Portability) support is enabled +- `engine.SUPPORT_MNNVL`: Whether MNNVL transport protocol support is enabled +- `engine.SUPPORT_MUSA`: Whether the Moore Threads MUSA GPU IPC transport is enabled +- `engine.SUPPORT_INTRA_NVLINK`: Whether intra-node NVLink support is enabled + +### Usage Example + +```python +from mooncake import engine + +# Check if CUDA is supported +if engine.SUPPORT_CUDA: + print("CUDA support is available") + +# Check all features +print(f"CUDA: {engine.SUPPORT_CUDA}") +print(f"EFA: {engine.SUPPORT_EFA}") +print(f"HIP: {engine.SUPPORT_HIP}") +print(f"MNNVL: {engine.SUPPORT_MNNVL}") +print(f"MUSA: {engine.SUPPORT_MUSA}") +print(f"Intra-NVLink: {engine.SUPPORT_INTRA_NVLINK}") +``` diff --git a/docs/source/api-reference/rust/index.md b/docs/source/api-reference/rust/index.md new file mode 100644 index 0000000000..59bb1b8f69 --- /dev/null +++ b/docs/source/api-reference/rust/index.md @@ -0,0 +1,15 @@ +# Rust + +| Module | Description | +|--------|-------------| +| [Mooncake Store](./mooncake-store) | Safe Rust bindings for the distributed KV-cache store (`mooncake_store`) | +| [Transfer Engine](./transfer-engine) | High-performance Rust library for Transfer Engine (`transfer_engine_rust`) | + +```{toctree} +:maxdepth: 1 +:hidden: + +mooncake-store +transfer-engine +``` + diff --git a/docs/source/api-reference/rust/mooncake-store.md b/docs/source/api-reference/rust/mooncake-store.md new file mode 100644 index 0000000000..f589246d9c --- /dev/null +++ b/docs/source/api-reference/rust/mooncake-store.md @@ -0,0 +1,223 @@ +# Mooncake Store Rust API + +This page documents the Rust crate `mooncake_store` (located at `mooncake-store/rust`). +It is a **safe wrapper** around the Mooncake Store C API (`store_c.h`). + +For deployment and service prerequisites, also see: + +- Mooncake Store deployment guide: `deployment/mooncake-store-deployment-guide` +- Error code reference: `troubleshooting/error-code` + +## Build & runtime prerequisites + +The Rust crate links against the C++ Mooncake build outputs. + +- **Build**: + - Build Mooncake with Store + Rust enabled via CMake: `-DWITH_STORE=ON -DWITH_STORE_RUST=ON` + - Or build with Cargo after exporting the CMake build directory / include paths (see `mooncake-store/rust/README.md`). +- **Runtime**: + - Dynamic linker must find Mooncake shared libraries (typically via `LD_LIBRARY_PATH` pointing at the CMake build outputs). + - The store client requires: + - a **metadata server** (HTTP metadata or etcd, depending on your deployment) + - `mooncake_master` + +## Quick start (copy-paste) + +```rust +use mooncake_store::MooncakeStore; + +fn main() -> Result<(), mooncake_store::StoreError> { + let store = MooncakeStore::new()?; + store.setup( + "127.0.0.1", + "http://127.0.0.1:8080/metadata", + 512 << 20, // global_segment_size + 128 << 20, // local_buffer_size + "tcp", + "", + "127.0.0.1:50051", + )?; + + store.put("hello", b"world", None)?; + let value = store.get("hello")?; + assert_eq!(value, b"world"); + + store.remove("hello", false)?; + Ok(()) +} +``` + +## API reference + +### `MooncakeStore` + +#### `new() -> Result` + +Allocate a new store handle (uninitialised). You must call `setup()` before any data operations. + +#### `setup(...) -> Result<(), StoreError>` + +Initialise the store client and establish connections. + +Parameters: + +- `local_hostname`: IP/hostname for this node. +- `metadata_server`: metadata URI, for example: + - HTTP: `"http://127.0.0.1:8080/metadata"` + - etcd: `"etcd://127.0.0.1:2379"` +- `global_segment_size`: per-segment size in bytes. +- `local_buffer_size`: local staging buffer size in bytes. +- `protocol`: transport protocol string (for example `"tcp"` / `"rdma"`). +- `device_name`: device selector; empty string means auto-select (when supported by the backend). +- `master_server_addr`: `mooncake_master` address, e.g. `"127.0.0.1:50051"`. + +Returns `Ok(())` on success, otherwise `StoreError::OperationFailed(code)`. + +#### `health_check() -> Result<(), StoreError>` + +Connectivity health check. Returns `Ok(())` when the backend is reachable. + +#### `put(key, value, config) -> Result<(), StoreError>` + +Store `value` under `key`. This is a **copying** API: `value` is copied into store-managed buffers. + +- `config`: optional replication settings (`ReplicateConfig`). + +#### `get(key) -> Result, StoreError>` + +Retrieve the full value for `key` into a newly allocated `Vec`. + +Notes: + +- Internally calls `get_size()` to allocate an exact-sized buffer, then `get_into()` to fill it. +- A missing key or backend failure can surface as `OperationFailed(...)` because the C API does not provide a distinct NotFound code in all paths. + +#### `unsafe get_into(key, buffer, size) -> Result` + +Retrieve the value for `key` into a caller-provided buffer. + +- **Returns**: number of bytes written on success. +- **Safety**: `buffer` must point to at least `size` bytes of writable valid memory. + +#### `is_exist(key) -> Result` + +Existence check. + +- `Ok(true)` if exists, `Ok(false)` if missing. +- Any other return code becomes `StoreError::OperationFailed(code)`. + +#### `get_size(key) -> Result` + +Get the stored value size in bytes. + +Important limitation: + +- The underlying C API uses a single negative return code for multiple error conditions, so Rust surfaces errors as `OperationFailed(raw_code)` without distinguishing NotFound. + +#### `get_hostname() -> Result` + +Returns the hostname (and potentially port) that the store client is registered under. + +#### `remove(key, force) -> Result<(), StoreError>` + +Remove a key. + +- If `force = true`, the key is removed even if another client is reading it. + +#### `remove_by_regex(pattern, force) -> Result` + +Remove keys matching a regex pattern. Returns number of removed keys. + +#### `remove_all(force) -> Result` + +Remove **all** keys. Returns number of removed keys. + +### Zero-copy APIs (advanced) + +The Rust wrapper exposes zero-copy APIs that map directly to the underlying RDMA-capable C++ store. + +#### `unsafe register_buffer(buffer, size) -> Result<(), StoreError>` + +Register a memory region for zero-copy operations. + +- **Safety**: `buffer` must remain valid and pinned until `unregister_buffer()` is called. +- This is required before calling `put_from()` or other registered-memory operations. + +#### `unsafe unregister_buffer(buffer) -> Result<(), StoreError>` + +Deregister a previously registered buffer. + +#### `unsafe put_from(key, buffer, size, config) -> Result<(), StoreError>` + +Store from a registered buffer. + +- **Safety**: `buffer` must have been registered via `register_buffer()` and be at least `size` bytes. + +### Batch APIs + +Batch forms are useful when you want to amortize RPC overhead. + +#### `unsafe batch_put_from(keys, buffers, sizes, config) -> Result, StoreError>` + +Batch version of `put_from()`. + +- **Returns**: per-key result codes (0 = success, non-zero = error code for that key). +- **Safety**: each `buffers[i]` must be registered and valid for `sizes[i]` bytes. + +#### `unsafe batch_get_into(keys, buffers, sizes) -> Result, StoreError>` + +Batch version of `get_into()`. + +- **Returns**: per-key bytes written (≥ 0) or error code (< 0). +- **Safety**: each destination buffer must be writable and at least `sizes[i]` bytes. + +#### `batch_is_exist(keys) -> Result, StoreError>` + +Batch existence check. Errors are returned as `OperationFailed(code)`. + +### `ReplicateConfig` + +Replication settings for write operations (`put`, `put_from`, `batch_put_from`). + +Fields: + +- `replica_num`: number of replicas (0 means “use server default”). +- `with_soft_pin`: prefer retaining the object in memory (soft pin). +- `with_hard_pin`: never evict (hard pin). +- `preferred_segments`: whitelist of segment names that should host a replica. + +Example: + +```rust +use mooncake_store::{MooncakeStore, ReplicateConfig}; + +fn write_with_replication(store: &MooncakeStore) -> Result<(), mooncake_store::StoreError> { + let cfg = ReplicateConfig { + replica_num: 2, + with_soft_pin: true, + with_hard_pin: false, + preferred_segments: vec!["seg-a".to_string(), "seg-b".to_string()], + }; + + store.put("k", b"v", Some(&cfg))?; + Ok(()) +} +``` + +### `StoreError` + +Errors returned by the Rust wrapper. + +- `NullHandle`: store handle allocation failed. +- `InvalidString`: input string contained an interior `\0` (cannot form C string). +- `OperationFailed(i32)`: underlying C layer returned a non-zero / negative code. +- `NotFound`: convenience for consumers that implement a NotFound check externally. +- `InvalidArgument(String)`: wrapper-level argument validation failure (e.g. mismatched array lengths). + +## Safety & thread-safety + +- `MooncakeStore` is `Send + Sync` (the underlying C object is internally synchronised). +- Methods that accept raw pointers are marked `unsafe`: + - You must uphold Rust aliasing and lifetime rules for buffers passed to FFI. + - For zero-copy operations, buffers must be registered and remain valid until unregistered. + diff --git a/docs/source/api-reference/rust/transfer-engine.md b/docs/source/api-reference/rust/transfer-engine.md new file mode 100644 index 0000000000..62ada671ec --- /dev/null +++ b/docs/source/api-reference/rust/transfer-engine.md @@ -0,0 +1,172 @@ +# Transfer Engine Rust API + +This page documents the Rust crate `transfer_engine_rust` (located at +`mooncake-transfer-engine/rust`). It is a **library** wrapper around the +Transfer Engine C API (`transfer_engine_c.h`). + +Hot-path types (`TransferRequest`, `BufferEntry`, `TransferStatus`) are +`#[repr(C)]` and compile-time layout-checked against the bindgen C types, so +`submit_transfer` passes a Rust slice to C with **no heap allocation and no +per-request copy** on the Rust side. + +For Transfer Engine design docs and non-Rust APIs, see: + +- Transfer Engine design docs: `design/transfer-engine/index` +- Transfer Engine C++ API: `api-reference/cpp/transfer-engine` +- Transfer Engine Python API: `api-reference/python/transfer-engine` + +## Build & runtime prerequisites + +- **Build**: + - Requires a Rust toolchain and libclang (bindgen). + - CMake: `-DWITH_RUST_EXAMPLE=ON`, then + `cmake --build build --target build_transfer_engine_rust`. + - Or Cargo after exporting `MOONCAKE_BUILD_DIR` / `MOONCAKE_TE_LIB_DIR` / + `MOONCAKE_TE_INCLUDE_DIR` (see `mooncake-transfer-engine/rust/README.md`). +- **Runtime**: + - Dynamic linker must find Transfer Engine shared libraries (`libasio.so`, …). + - A metadata server (HTTP metadata, etcd, or `P2PHANDSHAKE`) must be reachable. + - GitHub Actions runs `scripts/ci/run_transfer_engine_rust_smoke.sh` after the + C++ build (`cargo test --lib` plus the TCP loopback `minimal_smoke` test). + +## Quick start + +```rust +use transfer_engine_rust::{MemoryPool, TransferEngine, TransferRequest, WILDCARD_LOCATION}; + +fn main() -> Result<(), transfer_engine_rust::EngineError> { + let engine = TransferEngine::initialize( + "127.0.0.1:12345", + "http://127.0.0.1:8080/metadata", + "tcp", + "", + )?; + + let pool = MemoryPool::new(1 << 20); + unsafe { + engine.register_local_memory(pool.as_void_ptr(), pool.len(), WILDCARD_LOCATION)?; + + let seg = engine.open_segment("peer:12345")?; + let req = TransferRequest::write(pool.as_void_ptr(), seg, /*offset*/ 0, 4096); + engine.submit_and_wait(&[req], None)?; + + engine.unregister_local_memory(pool.as_void_ptr())?; + } + Ok(()) +} +``` + +`initialize(local_hostname, metadata_server, protocol, device_name)` matches +the Python constructor. `device_name` is accepted for API compatibility; NIC +filtering is done with `MC_TE_FILTERS` because the C ABI has no device-name +argument. + +The lower-level constructors map onto `createTransferEngine`: + +- `TransferEngine::new(metadata_uri, local_server_name, rpc_port)` +- `TransferEngine::create(TransferEngineOptions { … })` + +## Mental model + +- Register local memory regions as RDMA/TCP-capable buffers. +- Open a remote segment to obtain a `SegmentId`. +- Allocate a `BatchId` for a fixed number of requests. +- Submit a `&[TransferRequest]` (zero-copy FFI). +- Poll `get_transfer_status` / `wait_all`, then `free_batch_id`. + +Python-shaped helpers (`transfer_sync_write`, `batch_transfer_sync_*`, +`transfer_submit_write`) cache segment ids by hostname. They allocate a +batch internally. Use `submit_transfer` + `wait_all` when you need to keep +the batch/request arrays on the stack. + +## API reference + +### Types + +- `Opcode::{Read, Write}` — `OPCODE_READ` / `OPCODE_WRITE` +- `TransferStatusCode::{Waiting, Pending, Invalid, Canceled, Completed, Timeout, Failed}` +- `TransferRequest { opcode, source, target_id, target_offset, length }` — + layout matches `transfer_request_t`. Helpers: `TransferRequest::read`, + `TransferRequest::write`. +- `BufferEntry { addr, length }` — layout matches `buffer_entry_t` +- `TransferStatus { status, transferred_bytes }` — layout matches `transfer_status_t` +- `BatchId(u64)` — `INVALID_BATCH` on allocate failure +- `NotifyMsg { name, msg }` +- `NicLoadStat { device_name, inflight_bytes, ewma_bandwidth_bps }` +- `MemoryPool` — page-aligned, zeroed host buffer for registration +- `WILDCARD_LOCATION` (`"*"`), `LOCAL_SEGMENT` (`0`) + +### Engine lifecycle + +- `initialize(local_hostname, metadata_server, protocol, device_name)` +- `new` / `create` +- `discover_topology` +- `install_transport` / `uninstall_transport` +- `local_ip_and_port` +- Drop destroys the native handle (no double-free) + +### Memory + +All pointer APIs are `unsafe`. Registered memory must stay valid until +unregistered. + +- `register_local_memory` / `register_local_memory_ex` (remote-accessible flag) +- `unregister_local_memory` +- `register_memory` / `unregister_memory` — Python aliases using `WILDCARD_LOCATION` +- `register_local_memory_batch(&[BufferEntry])` — zero-copy +- `unregister_local_memory_batch(&[*mut c_void])` + +### Segments + +- `open_segment` / `open_segment_no_cache` / `open_segment_cached` +- `close_segment` +- `warmup_efa_segment` +- `remove_local_segment` +- `sync_segment_cache` + +### Transfers (zero-copy hot path) + +- `allocate_batch_id(batch_size)` +- `submit_transfer(batch_id, &[TransferRequest])` +- `submit_transfer_with_notify(batch_id, requests, &NotifyMsg)` +- `get_transfer_status(batch_id, task_id) -> TransferStatus` +- `wait_all(batch_id, count, timeout)` +- `submit_and_wait(&[TransferRequest], timeout)` — allocate + submit + wait + free +- `free_batch_id` + +### Python-shaped transfers + +These open (and cache) a segment by hostname: + +- `transfer_sync` / `transfer_sync_write` / `transfer_sync_read` +- `batch_transfer_sync` / `batch_transfer_sync_write` / `batch_transfer_sync_read` +- `transfer_submit_write` — returns `BatchId`; caller must `free_batch_id` +- `transfer_check_status` — polls task 0; does **not** free the batch + +### Notifications and diagnostics + +- `take_notifies() -> Vec` +- `send_notify(target_id, &NotifyMsg)` +- `nic_load_stats() -> Vec` +- `enable_graceful_shutdown` +- `show_links(json: bool) -> String` + +### Errors + +`EngineError` (`thiserror`, `#[non_exhaustive]`): + +- `NullHandle` +- `InvalidString` (interior NUL) +- `OperationFailed(i32)` — raw C status +- `InvalidArgument` +- `TransferFailed` +- `Timeout` + +## Safety & thread-safety + +- `TransferEngine` is `Send + Sync`; the C++ engine serializes internally. +- Pointer arguments must satisfy Rust aliasing and lifetime rules. +- Registered memory must remain valid until unregistered. +- `submit_transfer` does not copy request bytes; do not mutate a submitted + `TransferRequest` until the C call returns (the C layer copies into its own + vector before returning). diff --git a/docs/source/conf.py b/docs/source/conf.py index 817b116d9b..f339396447 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -41,35 +41,29 @@ "sphinx.ext.linkcode", "sphinx.ext.intersphinx", "sphinx_copybutton", - "autodoc2", "myst_parser", "sphinxarg.ext", "sphinx_design", "sphinx_togglebutton", + "sphinx_reredirects", + "sphinxcontrib.mermaid", ] myst_enable_extensions = [ "colon_fence", "fieldlist", ] -autodoc2_packages = [ - { - "path": "../../Mooncake", - "exclude_dirs": ["__pycache__", "third_party"], - }, -] -autodoc2_output_dir = "api" -autodoc2_render_plugin = "myst" -autodoc2_hidden_objects = ["dunder", "private", "inherited"] -autodoc2_sort_names = True -autodoc2_index_template = None - +myst_fence_as_directive = ["mermaid"] # Add any paths that contain templates here, relative to this directory. templates_path = ['_templates'] # List of patterns, relative to source directory, that match files and # directories to ignore when looking for source files. # This pattern also affects html_static_path and html_extra_path. -exclude_patterns: list[str] = ["**/*.template.md", "**/*.inc.md"] +exclude_patterns: list[str] = [ + "**/*.template.md", + "**/*.inc.md", + "zh_archive/**", +] # Exclude the prompt "$" when copying code copybutton_prompt_text = r"\$ " @@ -92,6 +86,7 @@ # Prevents the full API being added to the left sidebar of every page. # Reduces build time by 2.5x and reduces build size from ~225MB to ~95MB. 'collapse_navbar': True, + 'show_navbar_depth': 2, # Makes API visible in the right sidebar on API reference pages. 'show_toc_level': 3, } @@ -256,3 +251,95 @@ def linkcode_resolve(domain, info): } navigation_with_keys = False + +# Preserve published URLs when documentation is reorganized. Redirect targets +# are relative to the generated location of each legacy page. +redirects = { + "design/mooncake-store": "store/mooncake-store.html", + "design/ssd-offload": "store/ssd-offload.html", + "design/ssd-free-ratio-first-allocation": + "store/ssd-free-ratio-first-allocation.html", + "design/engram": "store/engram.html", + "design/unified-parallel-tensor-io": + "store/unified-parallel-tensor-io.html", + "design/tent/tebench": "../../performance/mooncake/tebench.html", + "deployment/ssd-offload": "ssd/ssd-offload.html", + "deployment/nvmf-ssd-deployment-guide": + "ssd/nvmf-ssd-deployment-guide.html", + "integrations/index": "../deployment/index.html", + "integrations/lmcache": + "../deployment/integrations/lmcache/index.html", + "integrations/lmdeploy": + "../deployment/integrations/lmdeploy.html", + "integrations/sglang/index": + "../../deployment/integrations/sglang/index.html", + "integrations/sglang/hicache-integration-v1": + "../../deployment/integrations/sglang/hicache-integration-v1.html", + "integrations/sglang/hicache-quick-start": + "../../deployment/integrations/sglang/hicache-quick-start.html", + "integrations/sglang/pd-disaggregation": + "../../deployment/integrations/sglang/pd-disaggregation.html", + "integrations/vllm/index": + "../../deployment/integrations/vllm/index.html", + "integrations/vllm/disagg-prefill-decode": + "../../deployment/integrations/vllm/disagg-prefill-decode.html", + "integrations/vllm/kv-cache-storage": + "../../deployment/integrations/vllm/kv-cache-storage.html", + "integrations/vllm/vllm-integration-v0.2": + "../../deployment/integrations/vllm/vllm-integration-v0.2.html", + "integrations/vllm/vllm-integration-v0.3": + "../../deployment/integrations/vllm/vllm-integration-v0.3.html", + "integrations/vllm/vllm-integration-v1.0": + "../../deployment/integrations/vllm/vllm-integration-v1.0.html", + "integrations/vllm/vllm-mooncakestoreconnector": + "../../deployment/integrations/vllm/vllm-mooncakestoreconnector.html", + "integrations/vllm/vllmv1-lmcache-integration": + "../../deployment/integrations/lmcache/vllmv1-lmcache-integration.html", + "integrations/vllm/vllmv1-lmcache-mp-integration": + "../../deployment/integrations/lmcache/vllmv1-lmcache-mp-integration.html", + "getting_started/examples/lmcache-integration": + "../../deployment/integrations/lmcache/index.html", + "getting_started/examples/lmdeploy-integration-v0.9": + "../../deployment/integrations/lmdeploy.html", + "getting_started/examples/sglang-integration-v1": + "../../deployment/integrations/sglang/pd-disaggregation.html", + "getting_started/examples/sglang-integration/index": + "../../../deployment/integrations/sglang/index.html", + "getting_started/examples/sglang-integration/hicache-integration-v1": + "../../../deployment/integrations/sglang/hicache-integration-v1.html", + "getting_started/examples/sglang-integration/hicache-quick-start": + "../../../deployment/integrations/sglang/hicache-quick-start.html", + "getting_started/examples/vllm-integration/index": + "../../../deployment/integrations/vllm/index.html", + "getting_started/examples/vllm-integration/disagg-prefill-decode": + "../../../deployment/integrations/vllm/disagg-prefill-decode.html", + "getting_started/examples/vllm-integration/kv-cache-storage": + "../../../deployment/integrations/vllm/kv-cache-storage.html", + "getting_started/examples/vllm-integration/vllm-integration-v0.2": + "../../../deployment/integrations/vllm/vllm-integration-v0.2.html", + "getting_started/examples/vllm-integration/vllm-integration-v0.3": + "../../../deployment/integrations/vllm/vllm-integration-v0.3.html", + "getting_started/examples/vllm-integration/vllm-integration-v1.0": + "../../../deployment/integrations/vllm/vllm-integration-v1.0.html", + "getting_started/examples/vllm-integration/vllm-mooncakestoreconnector": + "../../../deployment/integrations/vllm/vllm-mooncakestoreconnector.html", + "getting_started/examples/vllm-integration/vllmv1-lmcache-integration": + "../../../deployment/integrations/lmcache/vllmv1-lmcache-integration.html", + "getting_started/examples/vllm-integration/vllmv1-lmcache-mp-integration": + "../../../deployment/integrations/lmcache/vllmv1-lmcache-mp-integration.html", + "python-api-reference/dataproto-structured-object-transfer": + "../api-reference/python/dataproto-structured-object-transfer.html", + "python-api-reference/ep-backend": + "../api-reference/python/ep-backend.html", + "python-api-reference/mooncake-store": + "../api-reference/python/mooncake-store.html", + "python-api-reference/transfer-engine": + "../api-reference/python/transfer-engine.html", + "http-api-reference/http-service": + "../api-reference/http/http-service.html", + "design/transfer-engine/cpp-api": + "../../api-reference/cpp/transfer-engine.html", + "design/tent/cpp-api": "../../api-reference/cpp/tent.html", + "design/conductor/indexer-api-design": + "../../api-reference/http/conductor-indexer.html", +} diff --git a/docs/source/deployment/index.md b/docs/source/deployment/index.md new file mode 100644 index 0000000000..6e8a0cda0a --- /dev/null +++ b/docs/source/deployment/index.md @@ -0,0 +1,24 @@ +--- +orphan: true +--- + +# Deployment + +Deploy and operate Mooncake across standalone hosts, SSD-backed storage, and +Kubernetes environments. + +| Guide | Description | +|-------|-------------| +| [Mooncake Store Deployment and Tuning](mooncake-store-deployment-guide) | Configure Store clients, metadata services, storage tiers, and production tuning. | +| [KV Cache Sharing and Isolation](kv-cache-sharing-and-isolation) | Define cache-sharing boundaries across models, releases, request groups, and Mooncake tenants. | +| [Kubernetes Deployment](kubernetes-deployment-guide/index) | Deploy Mooncake Store and Transfer Engine integrations on Kubernetes. | +| [SSD Storage](ssd/index) | Configure local SSD offload or a shared NVMe-over-Fabrics storage pool. | + +## Framework Integrations + +| Integration | Description | +|-------------|-------------| +| [SGLang](integrations/sglang/index) | Deploy PD disaggregation and HiCache L3 storage with Mooncake. | +| [vLLM](integrations/vllm/index) | Deploy disaggregated prefill/decode and shared KV cache storage. | +| [LMCache](integrations/lmcache/index) | Use Mooncake as a distributed storage backend for LMCache. | +| [LMDeploy](integrations/lmdeploy) | Configure Mooncake as the PD disaggregation backend for LMDeploy. | diff --git a/docs/source/getting_started/examples/lmcache-integration.md b/docs/source/deployment/integrations/lmcache/index.md similarity index 96% rename from docs/source/getting_started/examples/lmcache-integration.md rename to docs/source/deployment/integrations/lmcache/index.md index 7dbcef4665..8723f4a8bc 100644 --- a/docs/source/getting_started/examples/lmcache-integration.md +++ b/docs/source/deployment/integrations/lmcache/index.md @@ -1,4 +1,6 @@ -# Mooncake x LMCache: Unite to Pioneer KVCache-Centric LLM Serving System +# Mooncake x LMCache Integration + +Mooncake and LMCache unite to pioneer KVCache-centric LLM serving systems. Mooncake and LMCache have announced a strategic collaboration aimed at pioneering a KVCache-centric Large Language Model (LLM) serving system. This partnership seeks to significantly enhance the efficiency, scalability, and responsiveness of LLM applications. @@ -23,7 +25,12 @@ This strategic alliance delivers substantial mutual technological advancements: For a complete deployment guide with step-by-step instructions, see: -👉 **[vLLM V1 Disaggregated Serving with Mooncake Store and LMCache](vllm-integration/vllmv1-lmcache-integration.md)** +:::{toctree} +:maxdepth: 1 + +vllmv1-lmcache-integration +vllmv1-lmcache-mp-integration +::: ## Performance Benchmarking and Results diff --git a/docs/source/deployment/integrations/lmcache/vllmv1-lmcache-integration.md b/docs/source/deployment/integrations/lmcache/vllmv1-lmcache-integration.md new file mode 100644 index 0000000000..df7fdd5961 --- /dev/null +++ b/docs/source/deployment/integrations/lmcache/vllmv1-lmcache-integration.md @@ -0,0 +1,191 @@ +# vLLM V1 Disaggregated Serving with Mooncake Store and LMCache + +## Overview + +This guide shows a two-machine 1-prefill/1-decode deployment using vLLM V1, +LMCache's non-MP `LMCacheConnectorV1`, and Mooncake Store as LMCache's remote +storage backend. + +LMCache supports both non-MP mode and MP mode with Mooncake Store. This page +covers the non-MP path, where each vLLM instance loads an LMCache YAML config +through `LMCACHE_CONFIG_FILE` and connects directly to Mooncake Store with +`remote_url: "mooncakestore://..."`. For the LMCache multiprocess server path, +see [vLLM V1 Disaggregated Serving with Mooncake Store and LMCache \[MP\]](vllmv1-lmcache-mp-integration.md). + +The examples below use: + +- Machine A: Mooncake master and vLLM decoder +- Machine B: vLLM prefiller +- Mooncake master RPC address: `{IP of Machine A}:50051` +- Mooncake HTTP metadata endpoint: `http://{IP of Machine A}:8080/metadata` +- RDMA device: `{RDMA device}` + +Replace these placeholders with the IP addresses, hostname, and RDMA device for +your environment. + +## Prerequisites + +Install Mooncake, vLLM, and LMCache on both machines. For installation details, +refer to the official documentation of each project: + +- [Mooncake build guide](../../../getting_started/build.md) +- [LMCache installation](https://docs.lmcache.ai/getting_started/installation.html) +- [vLLM installation](https://docs.vllm.ai/en/latest/getting_started/installation/) + +## Deployment + +### 1. Start Mooncake Master on Machine A + +```bash +mooncake_master -v=1 \ + --rpc_port=50051 \ + --metrics_port=9003 \ + --enable_http_metadata_server=true \ + --http_metadata_server_host=0.0.0.0 \ + --http_metadata_server_port=8080 +``` + +### 2. Configure and Start the vLLM Decoder on Machine A + +Modify the vLLM disaggregated prefill launcher to use a Mooncake-backed LMCache +config for the decoder. + +The decoder should continue to use `LMCacheConnectorV1` with `kv_role` set to +`kv_consumer`. + +```diff +diff --git a/examples/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh b/examples/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh +--- a/examples/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh ++++ b/examples/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh +@@ + elif [[ $1 == "decoder" ]]; then + # Decoder listens on port 8200 +- decode_config_file=$SCRIPT_DIR/configs/lmcache-decoder-config.yaml ++ decode_config_file=$SCRIPT_DIR/configs/mooncake-decoder-config.yaml +``` + +Create `configs/mooncake-decoder-config.yaml`: + +```yaml +chunk_size: 256 +remote_url: "mooncakestore://{IP of Machine A}:50051/" +remote_serde: "naive" +local_cpu: False +max_local_cpu_size: 100 + +extra_config: + local_hostname: "{IP of Machine A}" + metadata_server: "http://{IP of Machine A}:8080/metadata" + protocol: "rdma" + device_name: "{RDMA device}" + master_server_address: "{IP of Machine A}:50051" + global_segment_size: 32212254720 # 30GB + local_buffer_size: 1073741824 # 1GB + transfer_timeout: 1 + save_chunk_meta: False +``` + +Launch the decoder: + +```bash +bash disagg_vllm_launcher.sh decoder Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 +``` + +### 3. Configure and Start the vLLM Prefiller on Machine B + +Modify the launcher to use a Mooncake-backed LMCache config for the prefiller. + +The prefiller should continue to use `LMCacheConnectorV1` with `kv_role` set to +`kv_producer`. + +```diff +diff --git a/examples/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh b/examples/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh +--- a/examples/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh ++++ b/examples/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh +@@ + if [[ $1 == "prefiller" ]]; then + # Prefiller listens on port 8100 +- prefill_config_file=$SCRIPT_DIR/configs/lmcache-prefiller-config.yaml ++ prefill_config_file=$SCRIPT_DIR/configs/mooncake-prefiller-config.yaml +``` + +Create `configs/mooncake-prefiller-config.yaml`: + +```yaml +chunk_size: 256 +remote_url: "mooncakestore://{IP of Machine A}:50051/" +remote_serde: "naive" +local_cpu: False +max_local_cpu_size: 100 + +extra_config: + local_hostname: "{IP of Machine B}" + metadata_server: "http://{IP of Machine A}:8080/metadata" + protocol: "rdma" + device_name: "{RDMA device}" + master_server_address: "{IP of Machine A}:50051" + global_segment_size: 32212254720 # 30GB + local_buffer_size: 1073741824 # 1GB + transfer_timeout: 1 + save_chunk_meta: False +``` + +Launch the prefiller: + +```bash +bash disagg_vllm_launcher.sh prefiller Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 +``` + +### 4. Start the Disaggregated Proxy + +Use the LMCache [`disagg_proxy_server.py`](https://github.com/LMCache/LMCache/blob/dev/examples/disagg_prefill/disagg_proxy_server.py) to route requests between the prefiller and decoder. According to [LMCache/LMCache#1342](https://github.com/LMCache/LMCache/issues/1342), when using Mooncake Store as the backend, comment out the `wait_decode_kv_ready(...)` call in the proxy before starting it. + +```bash +python3 disagg_proxy_server.py \ + --host 0.0.0.0 \ + --port 9000 \ + --prefiller-host {IP of Machine B} \ + --prefiller-port 8100 \ + --decoder-host {IP of Machine A} \ + --decoder-port 8200 +``` + +### 5. Send a Test Request + +Send traffic to the proxy, not directly to either vLLM instance. + +```bash +curl -N http://{Proxy IP}:9000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4", + "messages": [ + { + "role": "user", + "content": "Explain how KV cache reuse helps long-context serving." + } + ], + "max_tokens": 128, + "temperature": 0.7 + }' +``` + +## Port and Configuration Checklist + +When changing ports away from these defaults, update all dependent settings +together: + +- Mooncake master `--rpc_port` must match `remote_url` and + `extra_config.master_server_address`. +- `extra_config.metadata_server` must point to the Mooncake HTTP metadata + endpoint when HTTP metadata is used. +- Decoder and prefiller `device_name`, `protocol`, `global_segment_size`, and + `local_buffer_size` should be set for the local hardware and workload. +- Proxy `--prefiller-port` and `--decoder-port` must match the two vLLM + instance ports. + +## Additional Resources + +* [Mooncake x LMCache: Unite to Pioneer KVCache-Centric LLM Serving System](index.md) +* [Using Mooncake in LMCache](https://docs.lmcache.ai/kv_cache/storage_backends/mooncake.html) +* [Using LMCache in vLLM](https://github.com/vllm-project/vllm/tree/main/examples/others/lmcache) diff --git a/docs/source/deployment/integrations/lmcache/vllmv1-lmcache-mp-integration.md b/docs/source/deployment/integrations/lmcache/vllmv1-lmcache-mp-integration.md new file mode 100644 index 0000000000..7af01e6ad4 --- /dev/null +++ b/docs/source/deployment/integrations/lmcache/vllmv1-lmcache-mp-integration.md @@ -0,0 +1,193 @@ +# vLLM V1 Disaggregated Serving with Mooncake Store and LMCache [MP] + +## Overview + +This guide shows a single-machine 1-prefill/1-decode deployment using vLLM V1, +LMCache's multiprocess server, and Mooncake Store as the LMCache L2 backend. + +LMCache supports both non-MP mode and MP mode with Mooncake Store. This page +covers the MP path, where vLLM instances connect to an LMCache server through +`LMCacheMPConnector`, and the LMCache server connects to Mooncake Store through +the `mooncake_store` L2 adapter. For the non-MP `LMCacheConnectorV1` path, see +[vLLM V1 Disaggregated Serving with Mooncake Store and LMCache](vllmv1-lmcache-integration.md). + +In this setup, one machine runs Mooncake master, one LMCache MP server, the +disaggregated proxy, the prefiller vLLM instance, and the decoder vLLM +instance. The prefiller and decoder should use different GPUs. + +This example uses `"metadata_server":"P2PHANDSHAKE"` for Mooncake transfer +metadata, so the Mooncake HTTP metadata server is not needed. If you switch to +HTTP metadata, remember that Mooncake's HTTP metadata server also defaults to +`8080`, which conflicts with LMCache's HTTP API on a single host. + +## Prerequisites + +Install Mooncake, vLLM, and LMCache on the machine. The example assumes an RDMA +deployment and uses: + +- local host address: `{IP of Machine}` +- RDMA device: `{RDMA device}` +- LMCache checkout path: `/path/to/LMCache` + +Replace these values with the local hostname/IP, RDMA device, and LMCache +checkout path for your environment. + +LMCache requirement: This example requires LMCache v0.4.5 or later. LMCache +must also be built from source with Mooncake support enabled, because the +`mooncake_store` MP L2 adapter depends on the optional +`lmcache.lmcache_mooncake` C++ extension. + +The standard prebuilt LMCache wheels currently include the Python Mooncake +adapter files, but do not include the optional `lmcache.lmcache_mooncake` +native extension. + +```bash +BUILD_MOONCAKE=1 \ +MOONCAKE_INCLUDE_DIR=/path/to/mooncake/include \ +MOONCAKE_LIB_DIR=/path/to/mooncake/lib \ +pip install -e /path/to/LMCache --verbose +``` + +## Deployment + +### 1. Start Mooncake Master + +```bash +mooncake_master -v=1 \ + --rpc_port=50051 \ + --metrics_port=9003 +``` + +### 2. Start the LMCache Multiprocess Server + +Start one LMCache MP server and configure Mooncake Store as the L2 adapter. + +```bash +lmcache server \ + --host 127.0.0.1 \ + --port 5555 \ + --http-host 127.0.0.1 \ + --http-port 8080 \ + --l1-size-gb 32 \ + --eviction-policy LRU \ + --no-l1-use-lazy \ + --l2-adapter '{ + "type": "mooncake_store", + "local_hostname": "{IP of Machine}", + "metadata_server": "P2PHANDSHAKE", + "protocol": "rdma", + "rdma_devices": "{RDMA device}", + "global_segment_size": 32212254720, + "local_buffer_size": 1073741824, + "master_server_addr": "127.0.0.1:50051" + }' +``` + +### 3. Start the Disaggregated Proxy + +The proxy receives client requests, sends prefill requests to the prefiller, +sends decode requests to the decoder, and receives LMCache request telemetry +from the prefiller. + +```bash +python /path/to/LMCache/examples/disagg_prefill_mp/disagg_proxy_server.py \ + --host 127.0.0.1 \ + --port 8000 \ + --prefiller-host 127.0.0.1 \ + --prefiller-port 8100 \ + --decoder-host 127.0.0.1 \ + --decoder-port 8200 \ + --telemetry-port 5768 +``` + +### 4. Start the vLLM Prefiller + +The prefiller reports request telemetry back to the proxy so the proxy knows +when KV cache storage has completed. + +```bash +CUDA_VISIBLE_DEVICES=0 \ +LMCACHE_REQUEST_TELEMETRY_TYPE=fastapi \ +LMCACHE_REQUEST_TELEMETRY_ENDPOINT=http://127.0.0.1:5768/api/v1/telemetry \ +vllm serve Qwen/Qwen3-4B \ + --host 127.0.0.1 \ + --port 8100 \ + --gpu-memory-utilization 0.8 \ + --no-enable-log-requests \ + --no-enable-prefix-caching \ + --kv-transfer-config '{ + "kv_connector": "LMCacheMPConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": { + "lmcache.mp.host": "tcp://127.0.0.1", + "lmcache.mp.port": 5555 + } + }' +``` + +### 5. Start the vLLM Decoder + +The decoder connects to the same local LMCache MP server. It does not need +request telemetry environment variables; only the prefiller reports the "KV +cache is stored" event back to the proxy. + +```bash +CUDA_VISIBLE_DEVICES=1 \ +vllm serve Qwen/Qwen3-4B \ + --host 127.0.0.1 \ + --port 8200 \ + --gpu-memory-utilization 0.8 \ + --no-enable-log-requests \ + --no-enable-prefix-caching \ + --kv-transfer-config '{ + "kv_connector": "LMCacheMPConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": { + "lmcache.mp.host": "tcp://127.0.0.1", + "lmcache.mp.port": 5555 + } + }' +``` + +### 6. Send a Test Request + +Send traffic to the proxy, not directly to either vLLM instance. + +```bash +curl -N http://127.0.0.1:8000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Qwen/Qwen3-4B", + "messages": [ + { + "role": "user", + "content": "Explain how KV cache reuse helps long-context serving." + } + ], + "max_tokens": 128, + "temperature": 0.7 + }' +``` + +## Port and Configuration Checklist + +When changing ports away from these defaults, update all dependent settings +together: + +- Mooncake master `--rpc_port` must match LMCache `master_server_addr`. +- The vLLM prefiller and decoder should connect to the local LMCache MP server + via `kv_connector_extra_config.lmcache.mp.host` and + `kv_connector_extra_config.lmcache.mp.port`. +- Proxy `--prefiller-port` and `--decoder-port` must match the two vLLM + `--port` values. +- `LMCACHE_REQUEST_TELEMETRY_ENDPOINT` on the prefiller must point to the proxy + telemetry endpoint. +- If `metadata_server` is changed from `P2PHANDSHAKE` to an HTTP metadata URL, + enable Mooncake HTTP metadata server and make sure its port does not conflict + with LMCache `--http-port`. + +## Additional Resources + +* [Mooncake x LMCache: Unite to Pioneer KVCache-Centric LLM Serving System](index.md) +* [LMCache MP `mooncake_store` L2 adapter](https://docs.lmcache.ai/mp/l2_storage.html#mooncake-store-mooncake-store-native-connector) +* [LMCache multiprocess disaggregated prefill example](https://github.com/LMCache/LMCache/tree/dev/examples/disagg_prefill_mp) diff --git a/docs/source/getting_started/examples/lmdeploy-integration-v0.9.md b/docs/source/deployment/integrations/lmdeploy.md similarity index 87% rename from docs/source/getting_started/examples/lmdeploy-integration-v0.9.md rename to docs/source/deployment/integrations/lmdeploy.md index 5de2a66994..50c9f9d7b5 100644 --- a/docs/source/getting_started/examples/lmdeploy-integration-v0.9.md +++ b/docs/source/deployment/integrations/lmdeploy.md @@ -16,17 +16,17 @@ pip install mooncake-transfer-engine Note: -- If any `.so` file is missing, uninstall the pip package with `pip3 uninstall mooncake-transfer-engine`, and build the binaries manually from source following the [build instructions](https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/build.md). +- If any `.so` file is missing, uninstall the pip package with `pip3 uninstall mooncake-transfer-engine`, and build the binaries manually from source following the [build instructions](../../getting_started/build.md). ### Install the latest version of LMDeploy -##### 1. Clone LMDeploy from the official repo +#### 1. Clone LMDeploy from the official repo ```bash git clone https://github.com/InternLM/lmdeploy.git ``` -##### 2. Build +#### 2. Build ##### 2.1 Build from source @@ -63,10 +63,10 @@ lmdeploy serve proxy \ ```bash lmdeploy serve api_server Qwen/Qwen3-8B \ - --server-name 192.168.0.101 \ - --server-port 23333 \ - --role Prefill \ - --proxy-url http://192.168.0.147:8000 \ + --server-name 192.168.0.101 \ + --server-port 23333 \ + --role Prefill \ + --proxy-url http://192.168.0.147:8000 \ --backend pytorch \ --migration-backend Mooncake ``` @@ -74,16 +74,16 @@ lmdeploy serve api_server Qwen/Qwen3-8B \ - The `--role` parameter sets the node role in the disaggregated system (`Prefill` for token embedding and KV cache generation). - The `--proxy-url` parameter connects the worker instance back to the proxy for coordination. - The `--backend` parameter specifies the model execution backend (e.g., `pytorch`, `turbomind`). -- The `--migration-backend` parameter defines the KV cache transport mechanism (e.g., `Mooncake` and `DlSllime`). +- The `--migration-backend` parameter defines the KV cache transport mechanism (e.g., `Mooncake` and `DLSlime`). #### Decode: ```bash lmdeploy serve api_server Qwen/Qwen3-8B \ - --server-name 192.168.0.147 \ - --server-port 23334 \ - --role Decode \ - --proxy-url http://192.168.0.147:8000 \ + --server-name 192.168.0.147 \ + --server-port 23334 \ + --role Decode \ + --proxy-url http://192.168.0.147:8000 \ --backend pytorch \ --migration-backend Mooncake ``` @@ -137,10 +137,10 @@ lmdeploy serve api_server Qwen/Qwen3-8B \ ```bash CUDA_VISIBLE_DEVICES=1 \ lmdeploy serve api_server Qwen/Qwen3-8B \ - --server-name 192.168.0.147 \ - --server-port 23334 \ - --role Decode \ - --proxy-url http://192.168.0.147:8000 \ + --server-name 192.168.0.147 \ + --server-port 23334 \ + --role Decode \ + --proxy-url http://192.168.0.147:8000 \ --backend pytorch \ --migration-backend Mooncake ``` @@ -164,4 +164,4 @@ curl -X POST "http://192.168.0.147:8000/v1/completions" \ - You can specify multiple prefill or decode instances with distinct `--server-port` and different GPUs using `CUDA_VISIBLE_DEVICES`. - MooncakeTransferEngine supports both intra-node (PCIe) and inter-node (RDMA/CXL) transfer, and device selection is automatic or customizable via config. - When using HF models that timeout during prefill, consider setting model path to `~/Qwen3-8B` to accelerate loading from localhost. -- Use `--log-level DEBUG` to get detailed runtime logs for troubleshooting. \ No newline at end of file +- Use `--log-level DEBUG` to get detailed runtime logs for troubleshooting. diff --git a/docs/source/getting_started/examples/sglang-integration/hicache-integration-v1.md b/docs/source/deployment/integrations/sglang/hicache-integration-v1.md similarity index 76% rename from docs/source/getting_started/examples/sglang-integration/hicache-integration-v1.md rename to docs/source/deployment/integrations/sglang/hicache-integration-v1.md index f47d48868f..7078faa02e 100644 --- a/docs/source/getting_started/examples/sglang-integration/hicache-integration-v1.md +++ b/docs/source/deployment/integrations/sglang/hicache-integration-v1.md @@ -6,22 +6,6 @@ This document describes how to use Mooncake as the storage backend for SGLang Hi ## Introduction -### About Mooncake - -Mooncake aims to enhance the inference efficiency of large language models (LLMs), especially in slow object storage environments, by constructing a multi-level caching pool on high-speed interconnected DRAM/SSD resources. Compared to traditional caching systems, Mooncake utilizes (GPUDirect) RDMA technology to transfer data directly in a zero-copy manner, while maximizing the use of multi-NIC resources on a single machine. - -For more details about Mooncake, please refer to [Mooncake project](https://github.com/kvcache-ai/Mooncake) and [Mooncake documents](https://kvcache-ai.github.io/Mooncake/). - -### About SGLang HiCache - -SGLang HiCache is a hierarchical KV caching system that extends SGLang's RadixAttention with advanced multi-tier memory management capabilities. It implements a scalable hierarchical storage architecture that spans GPU memory, CPU memory, and external storage layers, delivering significant performance improvements for large language model inference. - -HiCache introduces a **HiRadixTree** that acts as a page table for referencing KV caches across different memory tiers: **GPU Memory (L1)**, **CPU Memory (L2)**, **Mooncake and other Storage Backends (L3)**. - -The system includes an intelligent cache controller that automatically manages data movement between tiers, implementing optimized prefetching strategies and multiple write policies (write-through, write-through-selective, and write-back). - -For more details about SGLang HiCache, please refer to [HiCache system design document](https://docs.sglang.ai/advanced_features/hicache_design.html) and [this blog](https://lmsys.org/blog/2025-09-10-sglang-hicache/). - ### Mooncake & SGLang HiCache Mooncake serves as a high-performance L3 storage backend for SGLang HiCache, enabling distributed KV cache storage across multiple servers with RDMA-accelerated data transfer. This integration addresses the capacity limitations of traditional GPU-only or GPU+CPU caching by providing virtually unlimited cache storage through a distributed memory pool. @@ -42,21 +26,26 @@ This integration is particularly valuable for production deployments involving l ### Install SGLang -1. Clone SGLang from official repo +It is recommended to use uv for faster installation: ```bash -git clone git@github.com:sgl-project/sglang.git +pip install --upgrade pip +pip install uv +uv pip install sglang ``` -2. Build +The major version of Cuda is 13 by default. To install sglang under Cuda 12 with pip or uv, please try the following commands: ```bash -cd sglang pip install --upgrade pip -pip install -e "python[all]" +pip install uv +uv pip install sglang +uv pip install --force-reinstall torch==2.11.0 torchaudio==2.11.0 torchvision --index-url https://download.pytorch.org/whl/cu129 +uv pip install --force-reinstall sglang-kernel --index-url https://docs.sglang.ai/whl/cu129/ +uv pip install --force-reinstall sgl-deep-gemm --index-url https://docs.sglang.ai/whl/cu129/ --no-deps ``` -For more details, please refer to [SGLang official installation guide](https://docs.sglang.ai/get_started/install.html). +See the [SGLang official compilation guide](https://docs.sglang.ai/start/install.html) if you encounter issues. ### Install Mooncake @@ -66,37 +55,7 @@ For more details, please refer to [SGLang official installation guide](https://d pip install mooncake-transfer-engine ``` -**Method 2: from source** - -Clone Mooncake project: - -```bash -git clone https://github.com/kvcache-ai/Mooncake --recursive -``` - -Install dependencies: - -```bash -cd Mooncake -bash dependencies.sh -``` - -Build the project: - -```bash -mkdir build -cd build -cmake .. -make -j -``` - -Install Mooncake: - -```bash -sudo make install -``` - -For more details, please refer to [Mooncake official installation guide](https://kvcache-ai.github.io/Mooncake/getting_started/build.html). +If you want to build from source or using some other advanced features which not contained in prebuilt pip package, please refer to [Mooncake official installation guide](https://kvcache-ai.github.io/Mooncake/getting_started/build.html). ## Deployment @@ -138,7 +97,7 @@ mooncake_master --enable_http_metadata_server=true --http_metadata_server_port=8 When a `PutStart` request fails due to insufficient memory, or when the eviction thread detects that space usage has reached the configured high watermark ratio, an eviction task is triggered to free up space by evicting a portion of objects. -Due to memory fragmentation, allocation failures may occur even when memory usage has not yet reached 100%. The actual threshold depends on the workload. This [benchmark document](https://kvcache-ai.github.io/Mooncake/performance/allocator-benchmark-result.html) provides memory allocation efficiency results under different scenarios. if excessive allocation failures are observed, consider lowering this parameter accordingly. +Due to memory fragmentation, allocation failures may occur even when memory usage has not yet reached 100%. The actual threshold depends on the workload. This [benchmark document](https://kvcache-ai.github.io/Mooncake/performance/mooncake/allocator-benchmark-result.html) provides memory allocation efficiency results under different scenarios. if excessive allocation failures are observed, consider lowering this parameter accordingly. **Launch Mooncake `store service` (Optional):** @@ -365,69 +324,93 @@ python -m sglang.launch_server \ ### Prefill/Decode Disaggregation -In **PD disaggregation**, the configurations for the `metadata service`, `mooncake master`, and the optional `store service` remain the same as described above. The difference is that SGLang introduces three distinct roles: `prefill worker`, `decode worker`, and `router`. +Mooncake HiCache works with SGLang's **PD disaggregation** mode. The `master service`, `metadata service`, and optional `store service` configurations are the same as described above. + +1. Follow the [PD Disaggregation Guide](pd-disaggregation) to set up the prefill, decode, and router workers. +2. Add the HiCache-related parameters (`--enable-hierarchical-cache`, `--hicache-storage-backend mooncake`, `--hicache-storage-prefetch-policy`, etc.) to the **prefill worker** only, as described in the HiCache sections above. + +The Mooncake and HiCache configuration (environment variables or JSON config) is applied identically to the prefill worker — no changes are needed on the decode worker or router. + +## Troubleshooting + +**RDMA Registration Failure:** + +* In some environments, RDMA registration may require root privileges. In this case, try running the program as root. +* In certain environments (e.g., eRDMA), there is an upper limit on the total amount of RDMA memory that can be registered. Once this limit is exceeded, registration will fail. To resolve this, you can lower the value of `MOONCAKE_GLOBAL_SEGMENT_SIZE`, or reduce the host memory allocated to HiCache in the `SGLang server` (since this memory is fully registered with RDMA to enable zero-copy). + +**HiCache CPU Memory Usage:** -Among these, the `prefill worker` supports enabling **HiCache**. To run with PD disaggregation, start from the [PD configuration](https://kvcache-ai.github.io/Mooncake/getting_started/examples/sglang-integration-v1.html), and add the HiCache-related parameters (as previously described for the `SGLang server`) to the `prefill worker`. +When using HiCache, the default L2 host DRAM (CPU memory) size for KV cache is **2 times** the size of the L1 device memory (GPU memory) for KV cache. -In the example below, one `prefill worker`, one `decode worker`, and one `router` are launched. HiCache is enabled on the `prefill worker` to optimize prefill performance. +If the model is small but the GPU memory is large — especially in multi-TP (tensor parallel) setups — this may cause the L1 KV cache to become very large, which in turn can consume excessive CPU DRAM. -**Prefill worker**: +In such cases, you should manually configure an appropriate L2 cache size based on your hardware. This can be done by setting `--hicache-ratio` or `--hicache-size`. + +**HugeTLB Bring-up Checklist:** + +Before enabling HugeTLB-backed HiCache memory, size the host's hugepage pool against your planned `--hicache-size`, `MOONCAKE_GLOBAL_SEGMENT_SIZE`, and `MC_MMAP_ARENA_POOL_SIZE` values. + +From a source checkout: ```bash -MOONCAKE_TE_META_DATA_SERVER="http://127.0.0.1:8080/metadata" \ -MOONCAKE_MASTER=127.0.0.1:50051 \ -MOONCAKE_PROTOCOL="rdma" \ -MOONCAKE_DEVICE="mlx5_1" \ -MOONCAKE_GLOBAL_SEGMENT_SIZE=4294967296 \ -python -m sglang.launch_server \ - --model-path [model_path] \ - --page-size 64 \ - --enable-hierarchical-cache \ - --hicache-storage-prefetch-policy timeout \ - --hicache-storage-backend mooncake \ - --disaggregation-mode prefill \ - --disaggregation-ib-device "mlx5_1" \ - --base-gpu-id 0 \ - --port 30000 +python3 scripts/check_hicache_hugepage_requirements.py \ + --tp-size 4 \ + --hicache-size 64gb \ + --global-segment-size 8gb \ + --arena-pool-size 56gb \ + --available-hugetlb 512gb ``` -**Decode worker**: +From the source-built Docker image: ```bash -python -m sglang.launch_server \ - --model-path [model_path] \ - --page-size 64 \ - --disaggregation-mode decode \ - --disaggregation-ib-device "mlx5_1" \ - --base-gpu-id 1 \ - --port 30001 +mooncake-hicache-sizing \ + --tp-size 4 \ + --hicache-size 64gb \ + --global-segment-size 8gb \ + --arena-pool-size 56gb ``` -**Router**: +The `64gb` / `56gb` values above are tuned examples for large multi-GPU HiCache deployments, not defaults. Arena enablement remains opt-in, and the default pool is `8gb` only when you enable the arena via gflag without an explicit env override. On smaller hosts, start with `8gb` or `16gb` and size upward with the helper. + +The helper reports two numbers: + +* **Baseline floor** = `hicache-size + MOONCAKE_GLOBAL_SEGMENT_SIZE` per rank. Falling below this usually means startup or allocation failure. +* **Clean arena target** = baseline floor + `MC_MMAP_ARENA_POOL_SIZE` per rank. Reaching this target makes arena-backed launches less likely to spill onto the regular-page fallback path. + +Treat the helper as a planning tool rather than a hard proof: runs can still succeed between the baseline floor and the clean arena target if some arena allocations fall back to regular pages. + +Reserve HugeTLB pages on the host before launching SGLang: ```bash -python -m sglang_router.launch_router \ - --pd-disaggregation \ - --prefill "http://127.0.0.1:30000" \ - --decode "http://127.0.0.1:30001" \ - --host 0.0.0.0 \ - --port 8000 +sudo sysctl -w vm.nr_hugepages=262144 +grep -E 'HugePages_Total|HugePages_Free|Hugepagesize' /proc/meminfo ``` -## Troubleshooting +With `2 MiB` pages, `262144` pages equals `512 GiB`; `49152` pages equals `96 GiB`. Persist the setting with `/etc/sysctl.d/90-mooncake-hugepages.conf` if you need it across reboots. -**RDMA Registration Failure:** +**Memory Allocator Tuning:** -* In some environments, RDMA registration may require root privileges. In this case, try running the program as root. -* In certain environments (e.g., eRDMA), there is an upper limit on the total amount of RDMA memory that can be registered. Once this limit is exceeded, registration will fail. To resolve this, you can lower the value of `MOONCAKE_GLOBAL_SEGMENT_SIZE`, or reduce the host memory allocated to HiCache in the `SGLang server` (since this memory is fully registered with RDMA to enable zero-copy). +Mooncake's mmap arena is opt-in for HiCache host KV allocations. Setting `MC_MMAP_ARENA_POOL_SIZE` explicitly enables the arena and sizes the pool; the arena then pre-allocates a hugepage-backed pool and serves subsequent allocations via atomic bump pointer, reducing per-allocation latency from ~1000ns to ~50ns. The `56gb` example below is a benchmark-scale tuning value, not the allocator default. -**HiCache CPU Memory Usage:** +For HugeTLB-backed runs, export the hugepage and allocator settings together: -When using HiCache, the default L2 host DRAM (CPU memory) size for KV cache is **2 times** the size of the L1 device memory (GPU memory) for KV cache. +```bash +export MC_STORE_USE_HUGEPAGE=1 +export MC_STORE_HUGEPAGE_SIZE=2MB +export MOONCAKE_GLOBAL_SEGMENT_SIZE=8gb +export MC_MMAP_ARENA_POOL_SIZE=56gb +``` -If the model is small but the GPU memory is large — especially in multi-TP (tensor parallel) setups — this may cause the L1 KV cache to become very large, which in turn can consume excessive CPU DRAM. +To disable the arena and fall back to direct `mmap()` while keeping the hugepage-backed baseline path, set the flag before the first Mooncake mmap-buffer allocation in the process: -In such cases, you should manually configure an appropriate L2 cache size based on your hardware. This can be done by setting `--hicache-ratio` or `--hicache-size`. +```bash +export MC_DISABLE_MMAP_ARENA=1 +``` + +Without `MC_STORE_USE_HUGEPAGE=1`, the arena may opportunistically try hugepages and then retry on regular pages if HugeTLB is unavailable. If `MC_STORE_USE_HUGEPAGE=1` is set, Mooncake treats hugepages as a hard requirement for both the arena path and the direct-`mmap()` baseline path. It will not silently retry those host-buffer allocations on regular pages. + +If the helper reports `baseline_fits_arena_may_fallback`, either increase `vm.nr_hugepages` or reduce `MC_MMAP_ARENA_POOL_SIZE`. For containerized launches, pass the same environment variables through `docker run -e ...` and use `--ipc=host --ulimit memlock=-1 --shm-size=128g`. **More Information:** diff --git a/docs/source/getting_started/examples/sglang-integration/hicache-quick-start.md b/docs/source/deployment/integrations/sglang/hicache-quick-start.md similarity index 57% rename from docs/source/getting_started/examples/sglang-integration/hicache-quick-start.md rename to docs/source/deployment/integrations/sglang/hicache-quick-start.md index e234349b08..a4004f280f 100644 --- a/docs/source/getting_started/examples/sglang-integration/hicache-quick-start.md +++ b/docs/source/deployment/integrations/sglang/hicache-quick-start.md @@ -1,6 +1,6 @@ # Quick Start: SGLang HiCache with Mooncake Backend -Follow this streamlined workflow to get SGLang HiCache running with Mooncake as the L3 storage backend. +Follow this streamlined workflow to get SGLang HiCache running with Mooncake as the L3 storage backend. In benchmarks, pre-populated Mooncake achieves **best TTFT** across all tiers, maintaining high cache hit rates as conversation rounds grow ([details](../../../performance/sglang/sglang-hicache-benchmark-results-v1)). > Need more background or tuning options? See the [Complete Guide](hicache-integration-v1.md). @@ -10,21 +10,43 @@ Before you begin, make sure that: - SGLang is installed with HiCache support on the machine hosting your SGLang server. Refer to the [official installation guide](https://docs.sglang.ai/get_started/install.html) if needed. - Mooncake is installed and accessible as the hierarchical cache backend. Detailed build steps live in the [Mooncake documentation](https://kvcache-ai.github.io/Mooncake/getting_started/build.html). +- If you want HugeTLB-backed host KV memory, reserve the hugepages on the host before launching SGLang. The helper `python3 scripts/check_hicache_hugepage_requirements.py ...` from a source checkout, or `mooncake-hicache-sizing ...` inside the source-built Docker image, can size the required budget from your `--hicache-size`, `MOONCAKE_GLOBAL_SEGMENT_SIZE`, and `MC_MMAP_ARENA_POOL_SIZE` values. - The router package is installed to provide the `sglang_router` entrypoint: ```bash pip install sglang-router ``` -### 1. Launch the Mooncake master service +### 1. Size and reserve HugeTLB on the host + +```bash +python3 scripts/check_hicache_hugepage_requirements.py \ + --tp-size 4 \ + --hicache-size 64gb \ + --global-segment-size 8gb \ + --arena-pool-size 56gb \ + --available-hugetlb 512gb + +sudo sysctl -w vm.nr_hugepages=262144 +grep -E 'HugePages_Total|HugePages_Free|Hugepagesize' /proc/meminfo +``` + +`262144` is an example for `512 GiB` of `2 MiB` pages. The `64gb` / `56gb` sizing inputs above are tuned examples for large multi-GPU runs, not defaults. Adjust them, and the resulting `vm.nr_hugepages`, to match your planned HiCache budget. On smaller hosts, start with an arena pool of `8gb` or `16gb`. + +### 2. Launch the Mooncake master service ```bash mooncake_master --enable_http_metadata_server=true ``` -### 2. Launch SGLang with Mooncake L3 storage +### 3. Launch SGLang with Mooncake L3 storage ```bash +export MC_STORE_USE_HUGEPAGE="1" +export MC_STORE_HUGEPAGE_SIZE="2MB" +export MOONCAKE_GLOBAL_SEGMENT_SIZE="8gb" +export MC_MMAP_ARENA_POOL_SIZE="56gb" + MOONCAKE_MASTER=127.0.0.1:50051 python -m sglang.launch_server \ --model-path [model_path] \ --page-size 64 \ @@ -35,6 +57,8 @@ MOONCAKE_MASTER=127.0.0.1:50051 python -m sglang.launch_server \ **Key flag:** `--hicache-storage-prefetch-policy {best_effort,wait_complete,timeout}` determines when prefetching from storage should stop. `timeout` usually offers the best balance when Mooncake is the backend. +**Memory allocator:** Mooncake's mmap arena is opt-in. The `56gb` example above is a benchmark-scale tuning value, not the default. Setting `MC_MMAP_ARENA_POOL_SIZE` enables it and chooses the pool size; if you enable it via gflag instead, the default pool size is `8gb`. To force the baseline direct-`mmap()` path, set `MC_DISABLE_MMAP_ARENA=1` (also accepts `true`, `yes`, or `on`) before the first Mooncake mmap-buffer allocation in the process. See the [Complete Guide](hicache-integration-v1.md) for the full HugeTLB sizing and troubleshooting flow. + ## Prefill/Decode Disaggregation The disaggregated setup runs three processes—prefill worker, decode worker, and router. Launch each command below in its own terminal window. @@ -94,6 +118,7 @@ curl -X POST http://127.0.0.1:8000/generate \ - `--disaggregation-ib-device` is optional—SGLang autodetects devices, but you can set it explicitly (comma-separated, no spaces) when multiple NICs are available. - Use `--tp-size` to enable tensor-parallel execution across GPUs if required. +- For Docker launches, use `--ipc=host`, `--ulimit memlock=-1`, and a large `--shm-size`, then pass the same `MC_STORE_USE_HUGEPAGE`, `MC_STORE_HUGEPAGE_SIZE`, `MOONCAKE_GLOBAL_SEGMENT_SIZE`, and `MC_MMAP_ARENA_POOL_SIZE` values through `docker run -e ...`. - Optional flags to experiment with once you need them (some still have compatibility gaps): - `--disaggregation-decode-enable-offload-kvcache` writes the decode worker's outputs back into Mooncake; enable it when you want decoded KV to persist in L3. - Launch dedicated Mooncake `store service` nodes when you want to scale L3 capacity beyond what the SGLang servers contribute. diff --git a/docs/source/deployment/integrations/sglang/index.md b/docs/source/deployment/integrations/sglang/index.md new file mode 100644 index 0000000000..8926ad3c69 --- /dev/null +++ b/docs/source/deployment/integrations/sglang/index.md @@ -0,0 +1,54 @@ +# Mooncake x SGLang Integration + +Mooncake integrates with SGLang through two paths — **PD Disaggregation** for cross-instance KV cache transfer via the Transfer Engine, and **HiCache L3 Backend** for hierarchical KV cache storage with Mooncake Store. + +--- + +## PD Disaggregation + +SGLang uses Mooncake's Transfer Engine for direct zero-copy KV cache transfer between prefill and decode instances over RDMA, with support for EP and EPD backends. In benchmarks, PD disaggregation with Mooncake achieves **~30% lower ITL** while maintaining comparable throughput. + +``` + +-----------+ Transfer Engine (RDMA) +------------+ + | SGLang | ◄━━━━━━━━━━━━━━━━━━━━━━► | SGLang | + | Prefill | KV cache blocks | Decode | + +-----------+ +------------+ +``` + +**Related:** [Full PD Disaggregation Guide](pd-disaggregation) — installation, cross-node/same-node setup, XpYd topology, EP backend for MoE models, and EPD backend for multimodal models. + +**Benchmark:** [PD Disaggregation Performance](../../../performance/sglang/sglang-benchmark-results-v1) — compares 1P1D disaggregation with regular SGLang instances. + +--- + +## HiCache with Mooncake Store + +HiCache extends SGLang's RadixAttention with three memory tiers, using Mooncake Store as the distributed L3 backend. When local cache misses, HiCache automatically prefetches KV blocks from remote storage via RDMA. + +``` + +----------------------------------------------+ + | SGLang + HiCache | + | ┌─────────┐ ┌─────────┐ ┌──────────────┐ | + | │ L1(GPU) │ │ L2(CPU) │ │ L3(Mooncake) │ | + | └─────────┘ └─────────┘ └──────┬───────┘ | + +-----------------------------------+----------+ + | + +--------+---------+ + | Mooncake Store | + | Distributed Pool | + +------------------+ +``` + +**Related:** +- [HiCache Quick Start](hicache-quick-start) — minimal setup steps with hugepage sizing +- [HiCache Complete Guide](hicache-integration-v1) — full deployment, prefetch strategies, memory tuning, and architecture deep dive +- [SGLang Performance Benchmarks](../../../performance/sglang/index) — benchmark overview for PD disaggregation and HiCache with Mooncake + +::::{toctree} +:maxdepth: 1 +:hidden: + +pd-disaggregation +hicache-quick-start +hicache-integration-v1 +:::: diff --git a/docs/source/deployment/integrations/sglang/pd-disaggregation.md b/docs/source/deployment/integrations/sglang/pd-disaggregation.md new file mode 100644 index 0000000000..443c976b43 --- /dev/null +++ b/docs/source/deployment/integrations/sglang/pd-disaggregation.md @@ -0,0 +1,206 @@ +# SGLang Disaggregated Serving with MooncakeTransferEngine + +## Overview + +SGLang uses Mooncake's Transfer Engine to enable disaggregated prefill-decode (PD) serving across nodes via RDMA, with support for EP and EPD backends. This integration is based on [PR 4654](https://github.com/sgl-project/sglang/pull/4654) and [PR 4880](https://github.com/sgl-project/sglang/pull/4880). + +In benchmarks, PD disaggregation with Mooncake achieves **~30% lower ITL** while maintaining comparable throughput ([details](../../../performance/sglang/sglang-benchmark-results-v1)). + +``` + +-----------+ Transfer Engine (RDMA) +-----------+ + | SGLang | ◄━━━━━━━━━━━━━━━━━━━━━━► | SGLang | + | Prefill | KV cache blocks | Decode | + +-----------+ +-----------+ +``` + +## Prerequisites + +```bash +pip3 install mooncake-transfer-engine +``` + +If you encounter missing `lib*.so` errors, uninstall and build from source instead: + +```bash +pip3 uninstall mooncake-transfer-engine +# Build from source — see build instructions +``` + +### Install SGLang + +It is recommended to use uv for faster installation: + +```bash +pip install --upgrade pip +pip install uv +uv pip install sglang +``` + +The major version of Cuda is 13 by default. To install sglang under Cuda 12 with pip or uv, please try the following commands: + +```bash +pip install --upgrade pip +pip install uv +uv pip install sglang +uv pip install --force-reinstall torch==2.11.0 torchaudio==2.11.0 torchvision --index-url https://download.pytorch.org/whl/cu129 +uv pip install --force-reinstall sglang-kernel --index-url https://docs.sglang.ai/whl/cu129/ +uv pip install --force-reinstall sgl-deep-gemm --index-url https://docs.sglang.ai/whl/cu129/ --no-deps +``` + +See the [SGLang official compilation guide](https://docs.sglang.ai/start/install.html) if you encounter issues. + +## Configuration + +Key arguments: + +| Argument | Description | +|----------|-------------| +| `--disaggregation-mode` | `prefill` or `decode` — node role | +| `--disaggregation-ib-device` | RDMA device(s). Auto-detected, comma-separated for multi-NIC | +| `--tp-size` | Tensor parallelism size (optional) | +| `--base-gpu-id` | Starting GPU index for same-node deployments | +| `--host` / `--port` | SGLang service address | + +## Run PD Disaggregation + +### Multi-Node + +```bash +# Prefill node (192.168.0.137) +python -m sglang.launch_server \ + --model-path Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --disaggregation-mode prefill \ + --port 30000 --host 192.168.0.137 --tp-size 2 + +# Decode node (192.168.0.140) +python -m sglang.launch_server \ + --model-path Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --disaggregation-mode decode \ + --port 30001 --host 192.168.0.140 --tp-size 2 + +# Router +python3 -m sglang_router.launch_router \ + --pd-disaggregation \ + --prefill "http://192.168.0.137:30000" 8998 \ + --decode "http://192.168.0.140:30001" \ + --policy round_robin \ + --host 0.0.0.0 --port 8000 +``` + +Test: + +```bash +curl -X POST http://127.0.0.1:8000/generate \ + -H "Content-Type: application/json" \ + -d '{"text": "Tell me a long story", "sampling_params": {"temperature": 0}}' +``` + +### Single-Node + +```bash +# Prefill +python -m sglang.launch_server \ + --model-path Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --disaggregation-mode prefill \ + --port 30000 --host 192.168.0.137 --tp-size 2 + +# Decode (skip first 2 GPUs, use cards 2+) +python -m sglang.launch_server \ + --model-path Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --disaggregation-mode decode \ + --port 30001 --base-gpu-id 2 --host 192.168.0.137 --tp-size 2 + +# Router +python3 -m sglang_router.launch_router \ + --pd-disaggregation \ + --prefill "http://192.168.0.137:30000" 8998 \ + --decode "http://192.168.0.137:30001" \ + --policy round_robin \ + --host 0.0.0.0 --port 8000 + +# Router for multi node. +# Here is an example for 2 decode node running on 192.168.0.137 and 192.168.0.140 +python3 -m sglang_router.launch_router \ + --pd-disaggregation \ + --prefill "http://192.168.0.137:30000" 8998 \ + --decode "http://192.168.0.137:30001,http://192.168.0.140:30001" \ + --policy round_robin \ + --host 0.0.0.0 --port 8000 +``` + +TP is supported but not required — omit `--tp-size 2` for single-GPU setups. + +Multiple decode instances per prefill are supported. Multiple prefills on the same node are not supported due to bootstrap server port conflicts. + +```{tip} +If you encounter HuggingFace timeouts, set `export SGLANG_USE_MODELSCOPE=true`. +``` + +## Advanced Backends + +### EP Backend — Expert Parallelism for MoE Models + +For Mixture-of-Experts models (e.g., DeepSeek-V3), different nodes hold different expert weights. During inference, Mooncake transfers expert activations between prefill and decode nodes via RDMA, replacing the NCCL-based all-to-all with a faster, disaggregation-aware transfer path. + +```bash +# Prefill +python -m sglang.launch_server \ + --model-path deepseek-ai/DeepSeek-V3-0324 \ + --disaggregation-mode prefill \ + --port 30000 --host 192.168.0.137 \ + --tp-size 8 --dp-size 8 \ + --elastic-ep-backend mooncake \ + --moe-a2a-backend mooncake + +# Decode +python -m sglang.launch_server \ + --model-path deepseek-ai/DeepSeek-V3-0324 \ + --disaggregation-mode decode \ + --port 30001 --host 192.168.0.140 \ + --tp-size 8 --dp-size 8 \ + --elastic-ep-backend mooncake \ + --moe-a2a-backend mooncake +``` + +Set `--mooncake-ib-device` to the same value as `--disaggregation-ib-device` if explicit device specification is needed. + +### EPD Backend — Encoder-Prefill-Decode for Multimodal Models + +For multimodal models (e.g., LLaVA, InternVL), the **encoder** (Vision Transformer) processes images, the **prefill** node handles text + encoded visual tokens, and the **decode** node generates output. Mooncake transfers encoder outputs (visual embeddings) between these nodes via RDMA. + +The three node roles are: + +| Role | Flag | Responsibility | +|------|------|----------------| +| Encoder | `--encoder-only` | Processes images, produces visual embeddings. Mooncake sends embeddings to prefill via `--encoder-transfer-backend mooncake` | +| Prefill | `--disaggregation-mode prefill` | Receives encoder embeddings, runs prefill with text + visual context. Mooncake transfers KV cache to decode | +| Decode | `--disaggregation-mode decode` | Receives KV cache from prefill, generates output tokens | + +```bash +# Encoder-only node +python -m sglang.launch_server \ + --model-path $MODEL \ + --encoder-only \ + --encoder-transfer-backend mooncake \ + --port 30002 + +# Prefill node +python -m sglang.launch_server \ + --model-path $MODEL \ + --disaggregation-mode prefill \ + --disaggregation-transfer-backend mooncake \ + --encoder-transfer-backend mooncake \ + --tp $TP \ + --mem-fraction-static $MEM_FRACTION \ + --chunked-prefill-size $CHUNK_SIZE \ + --language-only \ + --encoder-urls http://127.0.0.1:30002 http://127.0.0.1:30003 \ + --port $PORT + +# Decode node +python -m sglang.launch_server \ + --model-path $MODEL \ + --disaggregation-mode decode \ + --disaggregation-transfer-backend mooncake \ + --port $PORT +``` diff --git a/docs/source/deployment/integrations/vllm/disagg-prefill-decode.md b/docs/source/deployment/integrations/vllm/disagg-prefill-decode.md new file mode 100644 index 0000000000..93453080f1 --- /dev/null +++ b/docs/source/deployment/integrations/vllm/disagg-prefill-decode.md @@ -0,0 +1,337 @@ +# Disaggregated Prefill-Decode with MooncakeConnector + +## Overview + +This guide demonstrates how to use `MooncakeConnector` with vLLM for disaggregated Prefill-Decode (PD) serving. `MooncakeConnector` enables direct cross-node KV cache transfer between prefill and decode instances using RDMA technology, achieving up to **142.25 GB/s** peak bandwidth (71.1% utilization of 8x RoCE). + +For more details about Mooncake, please refer to [Mooncake project](https://github.com/kvcache-ai/Mooncake) and [Mooncake documents](https://kvcache-ai.github.io/Mooncake/). + +--- + +## Choose Your vLLM Backend + +| Backend | vLLM Version | Status | Guide | +|---------|-------------|--------|-------| +| **vLLM V1** | Latest | Recommended | [Jump to V1 guide](#using-vllm-v1-recommended) | +| **vLLM V0** | ≤ v0.6.4.post1 | Legacy | [Jump to V0 guide](#using-vllm-v0-legacy) | + +```{admonition} New Users +:class: tip +If you are starting a new deployment, use the **vLLM V1** backend. V0 support is maintained for existing deployments only. +``` + +--- + +## Using vLLM V1 (Recommended) + +This section covers `MooncakeConnector` integration with vLLM V1 backend. The integration enables efficient cross-node KV cache transfer via RDMA. + +### Installation + +#### Prerequisites + +Install mooncake-transfer-engine through pip: + +```bash +pip install mooncake-transfer-engine +``` + +```{note} +If you encounter problems such as missing `lib*.so`, uninstall this package by `pip3 uninstall mooncake-transfer-engine`, and build the binaries manually according to the [instructions](../../../getting_started/build.md). +``` + +#### Install vLLM + +Refer to [vLLM official installation guide](https://docs.vllm.ai/en/latest/getting_started/installation.html) for the latest installation instructions. + +### Usage + +#### Basic Setup (Different Nodes) + +**Prefiller Node** (192.168.0.2): + +```bash +vllm serve Qwen/Qwen2.5-7B-Instruct \ + --port 8010 \ + --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_producer"}' +``` + +**Decoder Node** (192.168.0.3): + +```bash +vllm serve Qwen/Qwen2.5-7B-Instruct \ + --port 8020 \ + --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_consumer"}' +``` + +**Proxy Server:** + +```bash +# In vllm root directory. +python tests/v1/kv_connector/nixl_integration/toy_proxy_server.py \ + --prefiller-host 192.168.0.2 --prefiller-port 8010 \ + --decoder-host 192.168.0.3 --decoder-port 8020 +``` + +> NOTE: The Mooncake Connector currently uses the proxy from nixl_integration. This will be replaced with a self-developed proxy in the future. + +Now you can send requests to the proxy server through port 8000. + +**Test:** + +```bash +curl http://127.0.0.1:8000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Qwen/Qwen2.5-7B-Instruct", + "messages": [ + {"role": "user", "content": "Tell me a long story about artificial intelligence."} + ] + }' +``` + +#### Advanced Configuration + +**With Tensor Parallelism:** + +Prefiller: + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ +vllm serve Qwen/Qwen2.5-7B-Instruct \ + --port 8010 \ + --tensor-parallel-size 8 \ + --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_producer"}' +``` + +Decoder: + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ +vllm serve Qwen/Qwen2.5-7B-Instruct \ + --port 8020 \ + --tensor-parallel-size 8 \ + --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_consumer"}' +``` + +#### Configuration Parameters + +- `--kv-transfer-config`: JSON string to configure the KV transfer connector + - `kv_connector`: Set to `"MooncakeConnector"` + - `kv_role`: Role of the instance + - `kv_producer`: For prefiller instances that generate KV caches + - `kv_consumer`: For decoder instances that consume KV caches + - `kv_both`: Enables symmetric functionality (experimental) + - `num_workers`: Thread pool size in each prefiller worker to send kvcache (default 10) + +### Environment Variables + +- `VLLM_MOONCAKE_BOOTSTRAP_PORT`: Port for Mooncake bootstrap server (default: 8998) + - Required only for prefiller instances + - Each vLLM worker needs a unique port on its host + - For TP/DP deployments, each worker's port is computed as: `base_port + dp_rank * tp_size + tp_rank` +- `VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT`: Timeout (in seconds) for automatically releasing KV cache (default: 480) + - Used when a request is aborted to prevent holding resources indefinitely + +### Performance + +For detailed performance benchmarks and results, see the [vLLM PD Disaggregation Performance](../../../performance/vllm/vllm-v1-pd-performance.md) documentation. + +--- + +## Using vLLM V0 (Legacy) + +```{admonition} Legacy Backend +:class: warning +This section is for vLLM V0 backend (≤ v0.6.4.post1). For new deployments, use the [V1 backend](#using-vllm-v1-recommended) above. +``` + +This integration is based on [PR 10502](https://github.com/vllm-project/vllm/pull/10502) and [PR 10884](https://github.com/vllm-project/vllm/pull/10884). + +### Installation + +#### Prerequisite + +```bash +pip3 install mooncake-transfer-engine +``` + +```{note} +- If you encounter problems such as missing `lib*.so`, uninstall this package by `pip3 uninstall mooncake-transfer-engine`, and build the binaries manually according to the [instructions](../../../getting_started/build.md). +- For vLLM version ≤ v0.8.4, it requires `mooncake-transfer-engine ≤ 0.3.3.post2`. In the latest release, the interface `mooncake_vllm_adaptor` has been deprecated. +``` + +#### Install vLLM + +**1. Clone vLLM from official repo:** + +```bash +git clone git@github.com:vllm-project/vllm.git +``` + +**2. Build from source (Include C++ and CUDA code):** + +```bash +cd vllm +pip3 uninstall vllm -y +pip3 install -e . +``` + +```{tip} +If the build fails, try upgrading cmake: `pip3 install cmake --upgrade`. +``` + +If you encounter any problems, refer to the [vLLM official compilation guide](https://docs.vllm.ai/en/v0.6.4.post1/getting_started/installation.html#install-the-latest-code). + +### Configuration + +#### Prepare configuration file over RDMA + +Create a `mooncake.json` file for both Prefill and Decode instances. Use the identical config file on both sides. + +```json +{ + "prefill_url": "192.168.0.137:13003", + "decode_url": "192.168.0.139:13003", + "metadata_server": "192.168.0.139:2379", + "metadata_backend": "etcd", + "protocol": "rdma", + "device_name": "erdma_0" +} +``` + +- `prefill_url`: The IP address and port of the Prefill node (port is used to communicate with metadata server). +- `decode_url`: The IP address and port of the Decode node. If running prefill and decode on the same node, set a different port (at least 50 apart from `prefill_url` port) to avoid conflicts. +- `metadata_server`: The metadata server address. Supports `etcd`, `redis`, and `http` backends. Example: `"etcd://192.168.0.137:2379"`, `"redis://192.168.0.137:6379"`, `"http://192.168.0.137:8080/metadata"`. +- `metadata_backend`: Currently supports `"etcd"`, `"redis"`, and `"http"`. If absent and `metadata_server` has no prefix, defaults to `"etcd"`. This parameter will be deprecated in a future version. +- `protocol`: `"rdma"` or `"tcp"`. +- `device_name`: Required when protocol is `"rdma"`. Multiple NICs can be separated by commas (`"erdma_0,erdma_1"`). + +#### Prepare configuration file over TCP + +```json +{ + "prefill_url": "192.168.0.137:13003", + "decode_url": "192.168.0.139:13003", + "metadata_server": "192.168.0.139:2379", + "metadata_backend": "etcd", + "protocol": "tcp", + "device_name": "" +} +``` + +### Run Example + +Change the IP addresses and ports according to your environment. + +```bash +# Begin from root of your cloned repo! + +# 1. Start the etcd server +etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://localhost:2379 + +# 2. Run on the prefilling side (producer role) +MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_MODELSCOPE=True python3 -m vllm.entrypoints.openai.api_server \ + --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --port 8100 \ + --max-model-len 10000 \ + --gpu-memory-utilization 0.8 \ + --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2,"kv_buffer_size":2e9}' + +# 3. Run on the decoding side (consumer role) +MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_MODELSCOPE=True python3 -m vllm.entrypoints.openai.api_server \ + --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --port 8200 \ + --max-model-len 10000 \ + --gpu-memory-utilization 0.8 \ + --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2,"kv_buffer_size":2e9}' +``` + +**Key parameters:** +- `MOONCAKE_CONFIG_PATH`: Path to the mooncake.json configuration file. +- `VLLM_USE_MODELSCOPE`: Optional. Remove if you have HuggingFace access. +- `--kv-transfer-config`: Connector configuration + - `kv_connector`: `"MooncakeConnector"` + - `kv_role`: `"kv_producer"` or `"kv_consumer"` + - `kv_rank`: 0 for producer, 1 for consumer + - `kv_parallel_size`: Fixed to 2 currently + - `kv_buffer_size`: KVCache lookup buffer size; increase for longer prompts. If OOM occurs, decrease `--gpu-memory-utilization`. + - `kv_ip` and `kv_port`: Used to specify the IP address and port of the master node for `"PyNcclConnector"` distributed setup. Not used for `"MooncakeConnector"` currently. Instead, `"MooncakeConnector"` uses a config file to set up the distributed connection. +- `--tensor-parallel-size` / `-tp`: Supported. If running on the same node, set different `CUDA_VISIBLE_DEVICES`. + +```{note} +If running prefill and decode on the same node, set a different port for `decode_url`. To avoid port conflicts, ensure the decode port differs by at least 50 from the `prefill_url` port (e.g., `"decode_url": "192.168.0.137:13103"`). If the same URL is set for both, the port of `decode_url` will be automatically incremented by 100. +``` + +**Proxy Server:** + +```bash +python3 proxy_server.py +``` + +```python +# proxy_server.py +import os +import aiohttp +from quart import Quart, make_response, request + +AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60) +app = Quart(__name__) + +async def forward_request(url, data): + async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: + headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"} + async with session.post(url=url, json=data, headers=headers) as response: + if response.status == 200: + async for chunk_bytes in response.content.iter_chunked(1024): + yield chunk_bytes + +@app.route('/v1/completions', methods=['POST']) +async def handle_request(): + try: + original_request_data = await request.get_json() + prefill_request = original_request_data.copy() + prefill_request['max_tokens'] = 1 # prefill only + async for _ in forward_request('http://localhost:8100/v1/completions', prefill_request): + continue + generator = forward_request('http://192.168.0.139:8200/v1/completions', # Change IP + original_request_data) + response = await make_response(generator) + response.timeout = None + return response + except Exception as e: + import sys, traceback + exc_info = sys.exc_info() + print("Error occurred in disagg prefill proxy server") + print(e) + print("".join(traceback.format_exception(*exc_info))) + +if __name__ == '__main__': + app.run(host="0.0.0.0", port=8000) +``` + +> Be sure to change the IP address in the proxy server code. + +### Test + +```bash +curl -s http://localhost:8000/v1/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4", + "prompt": "San Francisco is a", + "max_tokens": 1000 + }' +``` + +--- + +## Troubleshooting + +- If you encounter connection issues, check that: + - All nodes can reach each other over the network + - Firewall rules allow traffic on the specified ports + - RDMA devices are properly configured and listed in `device_name` +- For missing library errors, rebuild `mooncake-transfer-engine` from source +- Enable debug logging with `VLLM_LOGGING_LEVEL=DEBUG` for detailed diagnostics +- For production deployments, consider using a more robust proxy solution diff --git a/docs/source/deployment/integrations/vllm/index.md b/docs/source/deployment/integrations/vllm/index.md new file mode 100644 index 0000000000..3c6d443c4f --- /dev/null +++ b/docs/source/deployment/integrations/vllm/index.md @@ -0,0 +1,55 @@ +# Mooncake x vLLM Integration + +## Overview + +Mooncake integrates with vLLM to accelerate large language model serving through high-performance KV cache transfer and shared storage. The integration supports two primary scenarios: + +- **Disaggregated Prefill-Decode Serving**: Seamlessly split prefill and decode across nodes using `MooncakeConnector`, with RDMA-powered cross-node KV cache transfer achieving up to **142.25 GB/s** peak bandwidth (71.1% utilization of 8x RoCE). Transfer overhead is negligible — for 32K-token prompts (4.50 GB of KV data), transfer takes only **31.65 ms**, accounting for just **4.2%** of total TTFT. +- **KV Cache Storage & Sharing**: Extend effective KV cache capacity via `MooncakeStore` / `MooncakeStoreConnector`, with hash-based prefix caching that enables multiple vLLM instances to share cached KV blocks. Supports CPU/Disk offloading and dynamic XpYd topologies at runtime. Distributed KV cache pool improves throughput by **3.8x**, reduces P50 TTFT and E2E latency by **46x** and **8.6x** (1P1D, 12GPUs), and scales to **60 GPUs** with >95% cache hit rate as shown in this [webpage](../../../performance/vllm/vllm-v1-mooncake-store.md). + +| Scenario | Guide | vLLM Backend | +|----------|-------|-------------| +| PD Disaggregation (KV transfer) | [Disaggregated Prefill-Decode](disagg-prefill-decode) | V1 ✅ / V0 ⚠️ | +| KV Cache Storage & Sharing | [KV Cache Storage with MooncakeStore](kv-cache-storage) | V1 ✅ / V0 ⚠️ | + +For detailed benchmark coverage across these scenarios, see [vLLM Integration Performance Benchmarks](../../../performance/vllm/index). + +```{admonition} New to Mooncake + vLLM? +:class: tip +Start with the V1 guides above. Legacy V0 documentation is available for existing deployments only. +``` + +--- + +## Getting Started + +### Disaggregated Prefill-Decode + +Direct KV cache transfer between prefill and decode nodes via `MooncakeConnector` using RDMA. + +::::{toctree} +:maxdepth: 1 + +disagg-prefill-decode +:::: + +### KV Cache Storage & Sharing + +Distributed KV cache storage via `MooncakeStore` / `MooncakeStoreConnector` for offloading, prefix caching, and cross-instance sharing. + +::::{toctree} +:maxdepth: 1 + +kv-cache-storage +:::: + +--- + +## Archived Documentation + +The following pages are from earlier versions of the integration and are no longer maintained. All content has been consolidated into the scenario-based guides above. + +- [MooncakeStoreConnector (Original)](vllm-mooncakestoreconnector) +- [vLLM V0 PD Disaggregation Demo (Original)](vllm-integration-v0.2) +- [vLLM V0 MooncakeStore (Original)](vllm-integration-v0.3) +- [vLLM V1 PD Disaggregation (Original)](vllm-integration-v1.0) diff --git a/docs/source/deployment/integrations/vllm/kv-cache-storage.md b/docs/source/deployment/integrations/vllm/kv-cache-storage.md new file mode 100644 index 0000000000..53b5f7a50c --- /dev/null +++ b/docs/source/deployment/integrations/vllm/kv-cache-storage.md @@ -0,0 +1,400 @@ +# KV Cache Storage & Sharing with MooncakeStore + +## Overview + +This guide demonstrates how to use `MooncakeStore` / `MooncakeStoreConnector` with vLLM to build a distributed KV cache storage pool. It enables KV cache offloading to CPU/SSD, hash-based prefix caching across multiple vLLM instances, and flexible XpYd disaggregated deployment — where you can dynamically adjust prefill and decode group sizes at runtime. + +Compared to Redis-based backends, MooncakeStore achieves significantly lower TTFT (e.g., **~32% improvement** in mean TTFT for 2P2D tp=2 under RDMA). + +--- + +## Choose Your vLLM Backend + +| Backend | Connector | vLLM Version | Status | Guide | +|---------|-----------|-------------|--------|-------| +| **vLLM V1** | `MooncakeStoreConnector` | Latest | Recommended | [Jump to V1 guide](#using-vllm-v1-recommended) | +| **vLLM V0** | `MooncakeStore` | ≤ v0.6.4.post1 | Legacy | [Jump to V0 guide](#using-vllm-v0-legacy) | + +```{admonition} New Users +:class: tip +If you are starting a new deployment, use the **vLLM V1** backend with `MooncakeStoreConnector`. V0 support is maintained for existing deployments only. +``` + +Key differences from v0.x to v1: +- **XpYd support and orchestration**: Dynamically change the population of prefill and decode groups +- **More stable and fault-tolerant**: A sudden crash of a single vLLM instance is tolerable; instance-to-instance connections are removed, so each instance works as a vanilla vLLM instance capable of handling requests independently + +--- + +## Using vLLM V1 (Recommended) + +This section covers `MooncakeStoreConnector` — the new vLLM KV connector that uses `MooncakeDistributedStore` as a shared KV cache pool. It enables: + +- **CPU/Disk offloading**: Extend effective KV cache capacity by offloading to CPU memory or SSD via Mooncake's transfer engine. +- **Hash-based prefix caching across instances**: Multiple vLLM instances share cached KV blocks through the store using block-hash deduplication. +- **Flexible deployment**: Works as a single-node KV cache extension (`kv_both`), or in disaggregated prefill-decode setups (`kv_producer` / `kv_consumer`). + +### Deployment + +#### 1. Prerequisites + +- [vLLM](https://github.com/vllm-project/vllm) is installed +- [Mooncake](https://github.com/kvcache-ai/Mooncake) is installed + +Refer to the [vLLM official repository](https://github.com/vllm-project/vllm) and [Mooncake official repository](https://github.com/kvcache-ai/Mooncake) for installation instructions and building from source. + +#### 2. Mooncake Master Server + +**Start:** + +```shell +mooncake_master --port 50063 +``` + +**Configure Mooncake**: Create a JSON configuration file (e.g., `mooncake_config.json`): + +```json +{ + "metadata_server": "http://127.0.0.1:8092/metadata", + "master_server_address": "127.0.0.1:50063", + "global_segment_size": "0", + "local_buffer_size": "2147483648", + "protocol": "rdma", + "device_name": "" +} +``` + +**Set environment variable:** + +```shell +export MOONCAKE_CONFIG_PATH=/path/to/mooncake_config.json +``` + +#### 3. Usage + +**3.1 Single-Node KV Cache Offloading** (`kv_both`): + +```shell +MOONCAKE_CONFIG_PATH=mooncake_config.json \ +vllm serve meta-llama/Llama-3.1-8B-Instruct \ + --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both"}' +``` + +**3.2 XpYd Disaggregated Prefill-Decode** (`kv_producer/kv_consumer`): + +Prefill Node: + +```shell +MOONCAKE_CONFIG_PATH=mooncake_config.json \ +VLLM_MOONCAKE_BOOTSTRAP_PORT=50052 \ +vllm serve meta-llama/Llama-3.1-8B-Instruct \ + --port 8100 \ + --kv-transfer-config '{ + "kv_connector": "MultiConnector", + "kv_role": "kv_producer", + "kv_connector_extra_config": { + "connectors": [ + { + "kv_connector": "MooncakeConnector", + "kv_role": "kv_producer" + }, + { + "kv_connector": "MooncakeStoreConnector", + "kv_role": "kv_producer" + } + ] + } + }' +``` + +Decode Node: + +```shell +MOONCAKE_CONFIG_PATH=mooncake_config.json \ +VLLM_MOONCAKE_BOOTSTRAP_PORT=50053 \ +vllm serve meta-llama/Llama-3.1-8B-Instruct \ + --port 8200 \ + --kv-transfer-config '{ + "kv_connector": "MultiConnector", + "kv_role": "kv_consumer", + "kv_connector_extra_config": { + "connectors": [ + { + "kv_connector": "MooncakeConnector", + "kv_role": "kv_consumer" + }, + { + "kv_connector": "MooncakeStoreConnector", + "kv_role": "kv_consumer" + } + ] + } + }' +``` + +Proxy: + +```shell +python examples/disaggregated/disaggregated_serving/mooncake_connector/mooncake_connector_proxy.py \ + --prefill http://192.168.0.2:8100 \ + --decode http://192.168.0.3:8200 +``` + +> When running with data parallelism, set a fixed `PYTHONHASHSEED` so that block hashes are consistent across DP ranks: +> +> ```shell +> PYTHONHASHSEED=0 vllm serve ... +> ``` +> +> Without this, identical prompts may produce different block hashes on different DP ranks, preventing cross-instance prefix cache hits. + +--- + +## Using vLLM V0 (Legacy) + +```{admonition} Legacy Backend +:class: warning +This section is for vLLM V0 backend with `MooncakeStore`. For new deployments, use the [V1 backend with `MooncakeStoreConnector`](#using-vllm-v1-recommended) above. +``` + +This integration is based on [PR 10502](https://github.com/vllm-project/vllm/pull/10502) and [PR 12957](https://github.com/vllm-project/vllm/pull/12957) to support KVCache transfer for intra-node and inter-node disaggregated serving. + +### Installation + +#### Prerequisite + +```bash +pip3 install mooncake-transfer-engine +``` + +```{note} +- If you encounter problems such as missing `lib*.so`, uninstall by `pip3 uninstall mooncake-transfer-engine`, and build manually according to the [instructions](../../../getting_started/build.md). +- For vLLM version ≤ v0.8.4, it requires `mooncake-transfer-engine ≤ 0.3.3.post2`. The interface `mooncake_vllm_adaptor` has been deprecated in the latest release. +``` + +#### Install vLLM + +**1. Clone vLLM:** + +```bash +git clone git@github.com:vllm-project/vllm.git +``` + +**2. Build from source:** + +```bash +cd vllm +pip3 install -e . +``` + +If you encounter problems, refer to the [vLLM official compilation guide](https://docs.vllm.ai/en/latest/getting_started/installation/index.html). + +### Configuration + +#### Prepare configuration for RDMA + +Create a `mooncake.json` file: + +```json +{ + "local_hostname": "192.168.0.137", + "metadata_server": "etcd://192.168.0.137:2379", + "protocol": "rdma", + "device_name": "erdma_0", + "master_server_address": "192.168.0.137:50001" +} +``` + +- `local_hostname`: The IP address of the current node. All prefill and decode instances on the same node can share this config. +- `metadata_server`: The metadata server. Supports `etcd`, `redis`, and `http` backends. +- `protocol`: `"rdma"` or `"tcp"`. +- `device_name`: Required for RDMA. Multiple NICs separated by commas (`"erdma_0,erdma_1"`). +- `master_server_address`: The IP address and port of the MooncakeStore master daemon. + +#### Prepare configuration for TCP + +```json +{ + "local_hostname": "192.168.0.137", + "metadata_server": "etcd://192.168.0.137:2379", + "protocol": "tcp", + "device_name": "", + "master_server_address": "192.168.0.137:50001" +} +``` + +### Run Example + +Change the IP addresses and ports according to your environment. `VLLM_USE_V1=0` is required for vLLM V0 backend. + +```bash +# Begin from root of your cloned repo! + +# 1. Start the etcd server +etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://localhost:2379 +# You may need to terminate other etcd processes before running the above command + +# 2. Start the mooncake_master server +mooncake_master --port 50001 +# If some vllm instances exit unexpectedly, some connection metadata will be +# corrupted since they are not properly cleaned. In that case, we recommend +# you restart the mooncake_master before running another test. + +# 3. Run multiple vllm instances +# kv_producer role +MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \ + --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --port 8100 \ + --max-model-len 10000 \ + --gpu-memory-utilization 0.8 \ + --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_producer"}' + +CUDA_VISIBLE_DEVICES=1 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \ + --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --port 8101 \ + --max-model-len 10000 \ + --gpu-memory-utilization 0.8 \ + --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_producer"}' + +CUDA_VISIBLE_DEVICES=2 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \ + --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --port 8102 \ + --max-model-len 10000 \ + --gpu-memory-utilization 0.8 \ + --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_producer"}' + +CUDA_VISIBLE_DEVICES=3 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \ + --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --port 8103 \ + --max-model-len 10000 \ + --gpu-memory-utilization 0.8 \ + --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_producer"}' + +# kv_consumer role +CUDA_VISIBLE_DEVICES=4 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \ + --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --port 8200 \ + --max-model-len 10000 \ + --gpu-memory-utilization 0.8 \ + --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer"}' + +CUDA_VISIBLE_DEVICES=5 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \ + --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --port 8201 \ + --max-model-len 10000 \ + --gpu-memory-utilization 0.8 \ + --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer"}' + +CUDA_VISIBLE_DEVICES=6 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \ + --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --port 8202 \ + --max-model-len 10000 \ + --gpu-memory-utilization 0.8 \ + --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer"}' + +CUDA_VISIBLE_DEVICES=7 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \ + --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --port 8203 \ + --max-model-len 10000 \ + --gpu-memory-utilization 0.8 \ + --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer"}' +``` + +**Key parameters:** +- `MOONCAKE_CONFIG_PATH`: Path to the mooncake.json configuration file. +- `VLLM_USE_MODELSCOPE`: Optional. Remove if you have HuggingFace access. +- `VLLM_USE_V1=0`: Required since the disaggregated feature is currently only supported on V0 vLLM. You can also `export` this configuration to the env instead of putting it in front of every command. +- `--model`: The model to use. +- `--port`: The vllm service port on which to listen. +- `--max-model-len`: The maximum length of the model. +- `--tensor-parallel-size` / `-tp`: Supported. All instances should have the same tensor_parallel_size. If running prefill and decode on the same node, set different `CUDA_VISIBLE_DEVICES` (e.g., `CUDA_VISIBLE_DEVICES=0,1` for prefill and `CUDA_VISIBLE_DEVICES=2,3` for decode). +- `--kv-transfer-config`: Set `kv_connector` to `"MooncakeStoreConnector"`, `kv_role` to `"kv_producer"`, `"kv_consumer"`, or `"kv_both"`. +- If some vLLM instances exit unexpectedly, connection metadata may be corrupted. Restart `mooncake_master` before another test. + +```bash +# 5. Start the proxy server +cd vllm +python3 examples/online_serving/disagg_examples/disagg_proxy_demo.py \ + --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --prefill localhost:8100 localhost:8101 \ + --decode localhost:8200 localhost:8201 \ + --port 8000 +``` + +- `--model`: The model and tokenizer used by the proxy server. +- `--port`: The proxy server port on which to listen. +- `--prefill` / `-p`: IP and port of the vllm prefill instances. +- `--decode` / `-d`: IP and port of the vllm decode instances. + +#### Dynamic XpYd Adjustment + +To dynamically adjust prefill and decode instances at runtime: + +```bash +export ADMIN_API_KEY="xxxxxxxx" + +# or add it before the command: +ADMIN_API_KEY="xxxxxxxx" python3 vllm/examples/online_serving/disagg_examples/disagg_demo.py \ + --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --prefill localhost:8100 localhost:8101 \ + --decode localhost:8200 localhost:8201 \ + --port 8000 \ + --scheduling round_robin + +# Add instances to groups dynamically +curl -X POST "http://localhost:8000/instances/add" \ + -H "Content-Type: application/json" \ + -H "X-API-Key: $ADMIN_API_KEY" \ + -d '{"type": "prefill", "instance": "localhost:8102"}' + +curl -X POST "http://localhost:8000/instances/add" \ + -H "Content-Type: application/json" \ + -H "X-API-Key: $ADMIN_API_KEY" \ + -d '{"type": "prefill", "instance": "localhost:8103"}' + +curl -X POST "http://localhost:8000/instances/add" \ + -H "Content-Type: application/json" \ + -H "X-API-Key: $ADMIN_API_KEY" \ + -d '{"type": "decode", "instance": "localhost:8202"}' + +curl -X POST "http://localhost:8000/instances/add" \ + -H "Content-Type: application/json" \ + -H "X-API-Key: $ADMIN_API_KEY" \ + -d '{"type": "decode", "instance": "localhost:8203"}' + +# Get proxy status +curl localhost:8000/status | jq +``` + +```{note} +Mooncake team provides this simple round-robin proxy as a demo. In production, you can implement custom global proxy strategies. +``` + +**Be sure to change the IP address in the commands.** + +### Test + +```bash +curl -s http://localhost:8000/v1/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4", + "prompt": "San Francisco is a", + "max_tokens": 1000 + }' +``` + +- If you are not testing on the proxy server, change `localhost` to the IP address of the proxy server. + +--- + +## Troubleshooting + +- If you encounter connection issues, check that: + - All nodes can reach each other over the network + - Firewall rules allow traffic on the specified ports + - RDMA devices are properly configured + - `mooncake_master` is running and reachable +- For missing library errors, rebuild `mooncake-transfer-engine` from source +- If vLLM instances exit unexpectedly, restart `mooncake_master` to clean up corrupted metadata +- Enable debug logging with `VLLM_LOGGING_LEVEL=DEBUG` for detailed diagnostics diff --git a/docs/source/getting_started/examples/vllm-integration/vllm-integration-v0.2.md b/docs/source/deployment/integrations/vllm/vllm-integration-v0.2.md similarity index 89% rename from docs/source/getting_started/examples/vllm-integration/vllm-integration-v0.2.md rename to docs/source/deployment/integrations/vllm/vllm-integration-v0.2.md index 1fc0ad176a..b40606594c 100644 --- a/docs/source/getting_started/examples/vllm-integration/vllm-integration-v0.2.md +++ b/docs/source/deployment/integrations/vllm/vllm-integration-v0.2.md @@ -1,11 +1,20 @@ +--- +orphan: true +--- + # vLLM V0 Disaggregated Serving Demo +```{admonition} Archived +:class: warning +This page has been **consolidated** into the unified [Disaggregated Prefill-Decode](disagg-prefill-decode) guide (see the V0 Legacy section). Please use that guide for up-to-date information. +``` + ## Overview -This is the latest version of mooncake-transfer-engine integration doc with the vLLM project based on [PR 10502](https://github.com/vllm-project/vllm/pull/10502) and [PR 10884](https://github.com/vllm-project/vllm/pull/10884) (vllm version: v0.6.4.post1/main) to accelerate KVCache transfer for inter-node disaggregated serving scenario. We have run some experiments to obtain some [preview benchmark results](../../../performance/vllm-benchmark-results-v0.2.md). More benchmark results will be released in due time. +This is the latest version of mooncake-transfer-engine integration doc with the vLLM project based on [PR 10502](https://github.com/vllm-project/vllm/pull/10502) and [PR 10884](https://github.com/vllm-project/vllm/pull/10884) (vllm version: v0.6.4.post1/main) to accelerate KVCache transfer for inter-node disaggregated serving scenario. **_Please note that this is still an experimental version and will be modified anytime based on feedback from the vLLM community._** - **Update(Apr 10, 2025)**: We are working on the vLLM v1 integration now. Stay tuned. - - **Update(Sep 5, 2025)**: We have released the vLLM v1 integration with Mooncake Store and LMCache. Please refer to [vllmv1-lmcache-integration](vllmv1-lmcache-integration.md) for more details. + - **Update(Sep 5, 2025)**: We have released the vLLM v1 integration with Mooncake Store and LMCache. Please refer to [vllmv1-lmcache-integration](../lmcache/vllmv1-lmcache-integration.md) for more details. ## Installation @@ -16,7 +25,7 @@ pip3 install mooncake-transfer-engine ``` Note: - - If you encounter problems such as missing `lib*.so`, you should uninstall this package by `pip3 uninstall mooncake-transfer-engine`, and build the binaries manually according to the [instructions](build.md). + - If you encounter problems such as missing `lib*.so`, you should uninstall this package by `pip3 uninstall mooncake-transfer-engine`, and build the binaries manually according to the [instructions](../../../getting_started/build.md). - For vLLM version <= v0.8.4, it requires mooncake-transfer-engine <= 0.3.3.post2. In the latest release, interface `mooncake_vllm_adaptor` has been deprecated. ### Install the latest version of vLLM @@ -89,10 +98,20 @@ etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://loc # You may need to terminate other etcd processes before running the above command # 2. Run on the prefilling side (producer role) -MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_MODELSCOPE=True python3 -m vllm.entrypoints.openai.api_server --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --port 8100 --max-model-len 10000 --gpu-memory-utilization 0.8 --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2,"kv_buffer_size":2e9}' +MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_MODELSCOPE=True python3 -m vllm.entrypoints.openai.api_server \ + --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --port 8100 \ + --max-model-len 10000 \ + --gpu-memory-utilization 0.8 \ + --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2,"kv_buffer_size":2e9}' # 3. Run on the decoding side (consumer role) -MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_MODELSCOPE=True python3 -m vllm.entrypoints.openai.api_server --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --port 8200 --max-model-len 10000 --gpu-memory-utilization 0.8 --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2,"kv_buffer_size":2e9}' +MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_MODELSCOPE=True python3 -m vllm.entrypoints.openai.api_server \ + --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --port 8200 \ + --max-model-len 10000 \ + --gpu-memory-utilization 0.8 \ + --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2,"kv_buffer_size":2e9}' ``` - `MOONCAKE_CONFIG_PATH` is the path to the mooncake.json configuration file. diff --git a/docs/source/getting_started/examples/vllm-integration/vllm-integration-v0.3.md b/docs/source/deployment/integrations/vllm/vllm-integration-v0.3.md similarity index 73% rename from docs/source/getting_started/examples/vllm-integration/vllm-integration-v0.3.md rename to docs/source/deployment/integrations/vllm/vllm-integration-v0.3.md index 821688653b..1e90e33123 100644 --- a/docs/source/getting_started/examples/vllm-integration/vllm-integration-v0.3.md +++ b/docs/source/deployment/integrations/vllm/vllm-integration-v0.3.md @@ -1,5 +1,14 @@ +--- +orphan: true +--- + # vLLM V0 Disaggregated Serving with MooncakeStore +```{admonition} Archived +:class: warning +This page has been **consolidated** into the unified [KV Cache Storage & Sharing](kv-cache-storage) guide (see the V0 Legacy section). Please use that guide for up-to-date information. +``` + ## Overview This is the latest version of the MooncakeStore integration doc with the vLLM project based on [PR 10502](https://github.com/vllm-project/vllm/pull/10502) and [PR 12957](https://github.com/vllm-project/vllm/pull/12957) to support KVCache transfer for intra-node and inter-node disaggregated serving scenario. Benchmark results will be released soon. @@ -13,7 +22,7 @@ Main changes from v0.x to v1: **_Please note that this is still an experimental version and will be modified anytime based on feedback from the vLLM community._** - **Update(Apr 10, 2025)**: We are working on the vLLM v1 integration now. Stay tuned. - - **Update(Sep 5, 2025)**: We have released the vLLM v1 integration with Mooncake Store and LMCache. Please refer to [vllmv1-lmcache-integration](vllmv1-lmcache-integration.md) for more details. + - **Update(Sep 5, 2025)**: We have released the vLLM v1 integration with Mooncake Store and LMCache. Please refer to [vLLM V1 LMCache integration](../lmcache/vllmv1-lmcache-integration.md) for more details. ## Installation @@ -24,7 +33,7 @@ pip3 install mooncake-transfer-engine ``` Note: - - If you encounter problems such as missing `lib*.so`, you should uninstall this package by `pip3 uninstall mooncake-transfer-engine`, and build the binaries manually according to the [instructions](build.md). + - If you encounter problems such as missing `lib*.so`, you should uninstall this package by `pip3 uninstall mooncake-transfer-engine`, and build the binaries manually according to the [instructions](../../../getting_started/build.md). - For vLLM version <= v0.8.4, it requires mooncake-transfer-engine <= 0.3.3.post2. In the latest release, interface `mooncake_vllm_adaptor` has been deprecated. ### Install the latest version of vLLM @@ -91,22 +100,62 @@ mooncake_master --port 50001 # 3. Run multiple vllm instances # kv_producer role -MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --port 8100 --max-model-len 10000 --gpu-memory-utilization 0.8 --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_producer"}' - -CUDA_VISIBLE_DEVICES=1 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --port 8101 --max-model-len 10000 --gpu-memory-utilization 0.8 --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_producer"}' - -CUDA_VISIBLE_DEVICES=2 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --port 8102 --max-model-len 10000 --gpu-memory-utilization 0.8 --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_producer"}' - -CUDA_VISIBLE_DEVICES=3 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --port 8103 --max-model-len 10000 --gpu-memory-utilization 0.8 --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_producer"}' +MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \ + --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --port 8100 \ + --max-model-len 10000 \ + --gpu-memory-utilization 0.8 \ + --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_producer"}' + +CUDA_VISIBLE_DEVICES=1 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \ + --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --port 8101 \ + --max-model-len 10000 \ + --gpu-memory-utilization 0.8 \ + --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_producer"}' + +CUDA_VISIBLE_DEVICES=2 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \ + --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --port 8102 \ + --max-model-len 10000 \ + --gpu-memory-utilization 0.8 \ + --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_producer"}' + +CUDA_VISIBLE_DEVICES=3 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \ + --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --port 8103 \ + --max-model-len 10000 \ + --gpu-memory-utilization 0.8 \ + --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_producer"}' # kv_consumer role -CUDA_VISIBLE_DEVICES=4 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --port 8200 --max-model-len 10000 --gpu-memory-utilization 0.8 --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer"}' - -CUDA_VISIBLE_DEVICES=5 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --port 8201 --max-model-len 10000 --gpu-memory-utilization 0.8 --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer"}' - -CUDA_VISIBLE_DEVICES=6 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --port 8202 --max-model-len 10000 --gpu-memory-utilization 0.8 --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer"}' - -CUDA_VISIBLE_DEVICES=7 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --port 8203 --max-model-len 10000 --gpu-memory-utilization 0.8 --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer"}' +CUDA_VISIBLE_DEVICES=4 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \ + --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --port 8200 \ + --max-model-len 10000 \ + --gpu-memory-utilization 0.8 \ + --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer"}' + +CUDA_VISIBLE_DEVICES=5 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \ + --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --port 8201 \ + --max-model-len 10000 \ + --gpu-memory-utilization 0.8 \ + --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer"}' + +CUDA_VISIBLE_DEVICES=6 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \ + --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --port 8202 \ + --max-model-len 10000 \ + --gpu-memory-utilization 0.8 \ + --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer"}' + +CUDA_VISIBLE_DEVICES=7 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python3 -m vllm.entrypoints.openai.api_server \ + --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --port 8203 \ + --max-model-len 10000 \ + --gpu-memory-utilization 0.8 \ + --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer"}' ``` - `MOONCAKE_CONFIG_PATH` is the path to the mooncake.json configuration file. @@ -127,7 +176,11 @@ CUDA_VISIBLE_DEVICES=7 MOONCAKE_CONFIG_PATH=./mooncake.json VLLM_USE_V1=0 python ```bash # 4. Start the proxy server cd vllm -python3 examples/online_serving/disagg_examples/disagg_proxy_demo.py --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --prefill localhost:8100 localhost:8101 --decode localhost:8200 localhost:8201 --port 8000 +python3 examples/online_serving/disagg_examples/disagg_proxy_demo.py \ + --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --prefill localhost:8100 localhost:8101 \ + --decode localhost:8200 localhost:8201 \ + --port 8000 ``` - The `--model` parameter specifies the model to use, also specifies the tokenizer used by the proxy server. @@ -139,7 +192,12 @@ python3 examples/online_serving/disagg_examples/disagg_proxy_demo.py --model Qwe # If you want to dynamically adjust the instances of p-nodes and d-nodes during runtime, you need to configure this environment variables. export ADMIN_API_KEY="xxxxxxxx" # or add it before the command: -ADMIN_API_KEY="xxxxxxxx" python3 vllm/examples/online_serving/disagg_examples/disagg_demo.py --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 --prefill localhost:8100 localhost:8101 --decode localhost:8200 localhost:8201 --port 8000 --scheduling round_robin +ADMIN_API_KEY="xxxxxxxx" python3 vllm/examples/online_serving/disagg_examples/disagg_demo.py \ + --model Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4 \ + --prefill localhost:8100 localhost:8101 \ + --decode localhost:8200 localhost:8201 \ + --port 8000 \ + --scheduling round_robin # Then use this command to add instances into prefill group or decode group curl -X POST "http://localhost:8000/instances/add" -H "Content-Type: application/json" -H "X-API-Key: $ADMIN_API_KEY" -d '{"type": "prefill", "instance": "localhost:8102"}' diff --git a/docs/source/getting_started/examples/vllm-integration/vllm-integration-v1.0.md b/docs/source/deployment/integrations/vllm/vllm-integration-v1.0.md similarity index 88% rename from docs/source/getting_started/examples/vllm-integration/vllm-integration-v1.0.md rename to docs/source/deployment/integrations/vllm/vllm-integration-v1.0.md index 3b140df6be..5cbc7de5d2 100644 --- a/docs/source/getting_started/examples/vllm-integration/vllm-integration-v1.0.md +++ b/docs/source/deployment/integrations/vllm/vllm-integration-v1.0.md @@ -1,5 +1,14 @@ +--- +orphan: true +--- + # vLLM v1 backend Disaggregated Serving with MooncakeConnector +```{admonition} Archived +:class: warning +This page has been **consolidated** into the unified [Disaggregated Prefill-Decode](disagg-prefill-decode) guide. Please use that guide for up-to-date information. +``` + ## Overview This guide demonstrates how to use the MooncakeConnector with vLLM v1 backend for disaggregated serving in Prefill-Decode separation architecture. The integration enables efficient cross-node KV cache transfer using RDMA technology. @@ -16,7 +25,7 @@ Install mooncake-transfer-engine through pip: pip install mooncake-transfer-engine ``` -Note: If you encounter problems such as missing `lib*.so`, you should uninstall this package by `pip3 uninstall mooncake-transfer-engine`, and build the binaries manually according to the [instructions](../build.md). +Note: If you encounter problems such as missing `lib*.so`, you should uninstall this package by `pip3 uninstall mooncake-transfer-engine`, and build the binaries manually according to the [instructions](../../../getting_started/build.md). ### Install vLLM @@ -45,7 +54,7 @@ vllm serve Qwen/Qwen2.5-7B-Instruct \ #### Proxy Server ```bash -# In vllm root directory. +# In vllm root directory. python tests/v1/kv_connector/nixl_integration/toy_proxy_server.py \ --prefiller-host 192.168.0.2 --prefiller-port 8010 \ --decoder-host 192.168.0.3 --decoder-port 8020 @@ -100,7 +109,7 @@ vllm serve Qwen/Qwen2.5-7B-Instruct \ - `kv_producer`: For prefiller instances that generate KV caches - `kv_consumer`: For decoder instances that consume KV caches - `kv_both`: Enables symmetric functionality (experimental) - - `num_workers`: Thread pool size in each prefiller worker to send kvcache (default 10) + - `num_workers`: Thread pool size in each prefiller worker to send kvcache (default 10) ## Environment Variables @@ -118,7 +127,7 @@ The following environment variables can be used to customize Mooncake behavior: ## Performance -For detailed performance benchmarks and results, see the [vLLM Benchmark](../../performance/vllm-v1-support-benchmark.md) documentation. +For detailed performance benchmarks and results, see the [vLLM PD Disaggregation Performance](../../../performance/vllm/vllm-v1-pd-performance.md) documentation. ## Notes diff --git a/docs/source/deployment/integrations/vllm/vllm-mooncakestoreconnector.md b/docs/source/deployment/integrations/vllm/vllm-mooncakestoreconnector.md new file mode 100644 index 0000000000..8517cb5de8 --- /dev/null +++ b/docs/source/deployment/integrations/vllm/vllm-mooncakestoreconnector.md @@ -0,0 +1,144 @@ +--- +orphan: true +--- + +# Guide: vLLM MooncakeStoreConnector + +```{admonition} Archived +:class: warning +This page has been **consolidated** into the unified [KV Cache Storage & Sharing](kv-cache-storage) guide (see the V1 Recommended section). Please use that guide for up-to-date information. +``` + +## Overview + +This document describes how to deploy vLLM's `MooncakeStoreConnector`. `MooncakeStoreConnector` is a new vLLM's KV connector that uses `MooncakeDistributedStore` as a shared KV cache pool. It enables: + +* **CPU/Disk offloading**: Extend effective KV cache capacity by offloading to CPU memory or SSD via Mooncake's transfer engine. +* **Hash-based prefix caching across instances**: Multiple vLLM instances share cached KV blocks through the store using block-hash deduplication. +* **Flexible deployment**: Works as a single-node KV cache extension (`kv_both`), or in disaggregated prefill-decode setups (`kv_producer` / `kv_consumer`). + + + +## Deployment + +### 1. Prerequisites + +Before you begin, make sure that: + +* [vLLM](https://github.com/vllm-project/vllm) is installed, [Mooncake](https://github.com/kvcache-ai/Mooncake) is installed. Refer to the [vLLM official repository](https://github.com/vllm-project/vllm) and [Mooncake official repository](https://github.com/kvcache-ai/Mooncake) for more installation instructions and building from source. + + + +### 2. Mooncake Master Server + +**Start:** + +```shell +mooncake_master --port 50063 +``` + +**Configure Mooncake** : Create a JSON configuration file (e.g., `mooncake_config.json`): + +```json +{ + "metadata_server": "http://127.0.0.1:8092/metadata", + "master_server_address": "127.0.0.1:50063", + "global_segment_size": "0", + "local_buffer_size": "2147483648", + "protocol": "rdma", + "device_name": "", +} +``` + +**Set environment variable:** + +```shell +export MOONCAKE_CONFIG_PATH=/path/to/mooncake_config.json +``` + + + +### 3. Usage + +**3.1** **Single-Node KV Cache Offloading** (i.e., `kv_both`) + +```shell +MOONCAKE_CONFIG_PATH=mooncake_config.json \ +vllm serve meta-llama/Llama-3.1-8B-Instruct \ + --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both"}' +``` + +**3.2** **XpYd Disaggregated Prefill-Decode** (i.e., `kv_producer/kv_consumer`) + +Prefill Node: + +```shell +MOONCAKE_CONFIG_PATH=mooncake_config.json \ +VLLM_MOONCAKE_BOOTSTRAP_PORT=50052 \ +vllm serve meta-llama/Llama-3.1-8B-Instruct \ + --port 8100 \ + --kv-transfer-config '{ + "kv_connector": "MultiConnector", + "kv_role": "kv_producer", + "kv_connector_extra_config": { + "connectors": [ + { + "kv_connector": "MooncakeConnector", + "kv_role": "kv_producer" + }, + { + "kv_connector": "MooncakeStoreConnector", + "kv_role": "kv_producer" + } + ] + } + }' +``` + +Decode Node: + +```shell +MOONCAKE_CONFIG_PATH=mooncake_config.json \ +VLLM_MOONCAKE_BOOTSTRAP_PORT=50053 \ +vllm serve meta-llama/Llama-3.1-8B-Instruct \ + --port 8200 \ + --kv-transfer-config '{ + "kv_connector": "MultiConnector", + "kv_role": "kv_consumer", + "kv_connector_extra_config": { + "connectors": [ + { + "kv_connector": "MooncakeConnector", + "kv_role": "kv_consumer" + }, + { + "kv_connector": "MooncakeStoreConnector", + "kv_role": "kv_consumer" + } + ] + } + }' +``` + +Proxy: + +```shell +python examples/disaggregated/disaggregated_serving/mooncake_connector/mooncake_connector_proxy.py \ + --prefill http://192.168.0.2:8100 \ + --decode http://192.168.0.3:8200 +``` + + + +> When running with data parallelism, set a fixed `PYTHONHASHSEED` so that block hashes are consistent across DP ranks: +> +> ```shell +> PYTHONHASHSEED=0 vllm serve ... +> ``` +> +> Without this, identical prompts may produce different block hashes on different DP ranks, preventing cross-instance prefix cache hits. + + +### 4. Performance + +Please refer to this [webpage](../../../performance/vllm/vllm-v1-mooncake-store.md). diff --git a/docs/source/deployment/kubernetes-deployment-guide/index.md b/docs/source/deployment/kubernetes-deployment-guide/index.md new file mode 100644 index 0000000000..315458d7ea --- /dev/null +++ b/docs/source/deployment/kubernetes-deployment-guide/index.md @@ -0,0 +1,47 @@ +# Kubernetes Deployment Guide + +Run Mooncake on Kubernetes as a shared **Store** cluster. The primary scenario below pairs it with SGLang **prefill/decode** inference — the Store serves as a HiCache L3 backend, while Mooncake's **Transfer Engine** moves KV cache directly between prefill and decode. The same Store also backs other orchestrators and engine stacks: see [RBG Integration](rbg-integration) for the RBG operator, and [llm-d Integration](llm-d-integration) for vLLM-based KV offloading and P/D transfer under llm-d. + +--- + +## Store + Transfer Engine (P/D disaggregation) + +A long-lived `mooncake-master` plus a replicated set of `mooncake-store` nodes form a shareable DRAM KV pool. The SGLang **prefill** pods use that pool as their hierarchical-cache L3 backend; **prefill and decode** use Mooncake's Transfer Engine for zero-copy P/D KV transfer over RDMA/TCP. A router fronts the prefill and decode endpoints. + +``` + Store cluster (no GPU) + +--------------------------------------------+ + | mooncake-master metadata + RPC | + | mooncake-store ×N (DRAM KV pool) | + +--------------------------------------------+ + ▲ + HiCache L3 │ (Get/Put) + metadata / RPC + │ + +-----┴-----+ Transfer +-----------+ + | SGLang | Engine | SGLang | + | Prefill |◄═════════════►| Decode | + | (GPU) | KV blocks | (GPU) | + +-----┬-----+ (RDMA/TCP) +-----┬-----+ + ▲ ▲ + │ │ + +--------+ +-----┴---------------------------┴-----+ + | client |──►| sglang-router | + +--------+ +---------------------------------------+ +``` + +**This section covers:** + +- [Mooncake on Kubernetes](mooncake-on-kubernetes) — stand up the Mooncake Store cluster with plain `Deployment` / `Service` objects. +- [RBG Integration](rbg-integration) — the full Store + P/D scenario with the [sgl-project/rbg](https://github.com/sgl-project/rbg) operator, including a production Mooncake cluster case. +- [llm-d Integration](llm-d-integration) — Mooncake's two integration points in [llm-d](https://github.com/llm-d/llm-d): the Store as a vLLM KV offload tier, and `MooncakeConnector` for P/D transfer. Links to the upstream llm-d examples. + +See also the [Mooncake Store Deployment & Tuning Guide](../mooncake-store-deployment-guide.md) for the component overview, client configuration, and tuning knobs. + +:::{toctree} +:maxdepth: 1 +:hidden: + +mooncake-on-kubernetes +rbg-integration +llm-d-integration +::: diff --git a/docs/source/deployment/kubernetes-deployment-guide/llm-d-integration.md b/docs/source/deployment/kubernetes-deployment-guide/llm-d-integration.md new file mode 100644 index 0000000000..e9f7a0ca46 --- /dev/null +++ b/docs/source/deployment/kubernetes-deployment-guide/llm-d-integration.md @@ -0,0 +1,76 @@ +# llm-d Integration + +[llm-d](https://github.com/llm-d/llm-d) packages vLLM inference on Kubernetes; Mooncake is the KV layer beneath it. The integration is a **three-layer stack**: + +- **Engine — vLLM.** The Mooncake connectors live here: `MooncakeStoreConnector` (offload) and `MooncakeConnector` (P/D transfer), driven by `--kv-transfer-config`, `mooncake_config.json`, and `PYTHONHASHSEED`. +- **Storage — Mooncake.** The Master/Client binaries, the Transfer Engine (RDMA), and the SSD tier. +- **Packaging — llm-d.** Kustomize overlays that wire the engine and storage into a Deployment + DaemonSet, plus a routing sidecar. + +Mooncake plugs in at **two independent points**, both configured on the vLLM side: + +- **KV cache offloading (storage)** — `MooncakeStoreConnector` uses a Mooncake Store pool as a shared offload tier. This is llm-d's primary Mooncake path. +- **P/D KV transfer** — the routing sidecar's `--kv-connector=mooncake` drives vLLM's `MooncakeConnector` (see [below](#pd-transfer)). + +The two share the Transfer Engine but serve different purposes; vLLM's `MultiConnector` can compose them when a deployment needs both. For each layer, this page points at the authoritative source and keeps only what that source leaves for the integrator to reconcile. + +## vLLM connector contract + +The vLLM side of the integration is a single command. With the Mooncake Master already running and a `mooncake_config.json` on disk, `MooncakeStoreConnector` is wired on a single node with that config file and a connector spec: + +```bash +MOONCAKE_CONFIG_PATH=mooncake_config.json \ +vllm serve meta-llama/Llama-3.1-8B-Instruct \ + --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both"}' +``` + +The `mooncake_config.json` this command names has a `mode` field that selects the storage topology: + +| Mode | Who owns the DRAM pool | `global_segment_size` | Components | +|---|---|---|---| +| `embedded` | each vLLM rank contributes DRAM in-process | **> 0** (e.g. `"80GB"`) | Master + vLLM | +| `standalone-store` | an external `mooncake_client` owns the CPU DRAM + SSD pool; vLLM ranks are pure requesters | **`0`** | Master + Client + vLLM | + +The connector enforces the pairing at startup: `embedded` requires `global_segment_size > 0`, `standalone-store` requires `global_segment_size == 0`, and `local_buffer_size` must be `> 0` in both. `standalone-store` decouples the pool from vLLM — it survives vLLM restarts and can live on GPU-less nodes with large DRAM and NVMe. The SSD tier is a separate switch, not implied by the mode: `enable_offload` must be set together on the vLLM config, the Master, and the Client (vLLM even accepts `embedded` with `enable_offload`). + +llm-d instantiates this same contract at scale: the `--kv-transfer-config`, the `MOONCAKE_CONFIG_PATH` JSON, and `PYTHONHASHSEED` set here are what the guide's `patch-vllm.yaml` sets on each model-server Pod. vLLM's docs are the authoritative field reference: + +- [MooncakeStore connector usage](https://github.com/vllm-project/vllm/blob/main/docs/features/mooncake_store_connector_usage.md) — the `mooncake_config.json` schema (including the `mode` field), the single-node form above, and the `MultiConnector[MooncakeConnector + MooncakeStoreConnector]` composition for XpYd. +- [Mooncake connector usage](https://github.com/vllm-project/vllm/blob/main/docs/features/mooncake_connector_usage.md) — the P/D transfer connector and its bootstrap port (`8998`). + +```{important} +Set `PYTHONHASHSEED` to the **same fixed value on every instance sharing a Store**. Mooncake keys blocks by a content hash derived from vLLM's block hashes, and Python randomises its hash seed per process by default — leave it unset and two instances compute different keys for identical tokens, so the pool stays healthy while the cross-instance hit rate sits at zero. +``` + +## KV cache offloading on K8s (llm-d) + +llm-d's [Tiered Prefix Cache guide](https://github.com/llm-d/llm-d/tree/v0.8.0/guides/tiered-prefix-cache) packages MooncakeStore as an offload backend: evicted KV blocks move out of GPU HBM into a Mooncake Store pool that serves them back on a later hit, saving the recompute. Because the pool is shared, multiple vLLM instances reuse each other's cached prefixes. The guide carries the deployment sequence, node sizing, and manifests; deploy it with `helm install` for the router chart and `kubectl apply -k` over the Mooncake overlays. + +```{note} +**Upstream sources** (llm-d `v0.8.0`): +- Guide: [Tiered Prefix Cache](https://github.com/llm-d/llm-d/tree/v0.8.0/guides/tiered-prefix-cache) — the deployment walkthrough; the MooncakeStore path has `cpu` and `fs` variants. +- Architecture: [KV Offloader](https://github.com/llm-d/llm-d/blob/v0.8.0/docs/architecture/advanced/kv-management/kv-offloader.md) — Master/Client roles and a `mooncake_config.json` field table. +- Overlays: [`helpers/mooncake-master-store/`](https://github.com/llm-d/llm-d/tree/v0.8.0/helpers/mooncake-master-store) (Master; both variants), [`modelserver/gpu/vllm/mooncake-store/`](https://github.com/llm-d/llm-d/tree/v0.8.0/guides/tiered-prefix-cache/modelserver/gpu/vllm/mooncake-store) (`cpu`/`fs` vLLM overlays), [`helpers/mooncake-client/`](https://github.com/llm-d/llm-d/tree/v0.8.0/helpers/mooncake-client) (Client DaemonSet base; `fs` patches it up). +``` + +```{note} +**Confirming Mooncake specifically.** The guide's offload verification inspects the decode Pod's `/mnt/files-storage/kv-cache`, which is the vLLM-native/LMCache filesystem tier — a Mooncake `fs` deployment writes nowhere near it, so an all-`Ready` cluster does not prove the Store is working. Confirm Mooncake on the `mooncake_client` Pod instead: its `/data/mooncake-offload` should grow after a long request, and the Master should show the pool registered (metrics on `:9003`). +``` + +(pd-transfer)= +## P/D KV transfer (MooncakeConnector) + +The P/D transfer path is independent of the Store. vLLM ships a runnable single-node example for it — [`examples/disaggregated/mooncake_connector/`](https://github.com/vllm-project/vllm/tree/main/examples/disaggregated/mooncake_connector) (proxy + launch scripts, configurable model and bootstrap port). In a prefill/decode disaggregated deployment, llm-d's routing sidecar drives the same connector: + +```bash +# on the decode pod's routing sidecar +--kv-connector=mooncake +``` + +The sidecar queries a bootstrap endpoint on the prefill pods to resolve the target engine, then dispatches prefill and decode so the decoder pulls KV directly from the prefiller. + +| Setting | Default | Notes | +|---|---|---| +| `--kv-connector=mooncake` | — | Selects vLLM's `MooncakeConnector` on the sidecar. The router accepts only fixed connector names (`mooncake`, not `MultiConnector`), so the serving pods must run `MooncakeConnector` to match it — either as the top-level `kv_connector`, or nested inside a `MultiConnector`. | +| `--mooncake-bootstrap-port` / `MOONCAKE_BOOTSTRAP_PORT` | `8998` | Port of the Mooncake bootstrap endpoint on prefill pods. Corresponds to vLLM's `VLLM_MOONCAKE_BOOTSTRAP_PORT`. | + +**Capability vs. packaged path.** The `mooncake` connector reached a released sidecar in **llm-d-router v0.9.0** ([`supportedKVConnectors` at that tag](https://github.com/llm-d/llm-d-router/blob/v0.9.0/pkg/sidecar/proxy/options.go)), and the umbrella [llm-d v0.8.0](https://github.com/llm-d/llm-d/tree/v0.8.0) release pins that sidecar — so the sidecar **capability** is present. The **deployment path** is reference-only: neither v0.8.0 nor `main` ships a Mooncake P/D overlay, and llm-d's packaged [P/D disaggregation guide](https://github.com/llm-d/llm-d/tree/main/guides/pd-disaggregation) still wires `--kv-connector=nixlv2`, so standing up Mooncake P/D means adapting the serving-pod and sidecar manifests by hand. For the current connector table and flags, see the sidecar's [disaggregation reference](https://github.com/llm-d/llm-d-router/blob/main/docs/disaggregation.md). diff --git a/docs/source/deployment/kubernetes-deployment-guide/mooncake-on-kubernetes.md b/docs/source/deployment/kubernetes-deployment-guide/mooncake-on-kubernetes.md new file mode 100644 index 0000000000..8316706856 --- /dev/null +++ b/docs/source/deployment/kubernetes-deployment-guide/mooncake-on-kubernetes.md @@ -0,0 +1,134 @@ +# Mooncake on Kubernetes + +Deploy a Mooncake Store cluster — a `mooncake-master` plus replicated `mooncake-store` nodes — with plain Kubernetes objects (`Deployment` and `Service`). + +Use it together with the [Mooncake Store Deployment & Tuning Guide](../mooncake-store-deployment-guide.md): that guide explains the components and tuning knobs; this page maps them to Kubernetes objects. + +## Deploy the Mooncake Store cluster + +A shareable Store cluster has one `mooncake-master` (the RPC coordinator) and a replicated set of stateless `mooncake-store` nodes that contribute DRAM to the pool. The nodes use Mooncake's P2P handshake (`P2PHANDSHAKE`) for Transfer Engine peer discovery, so there is no separate metadata service to run — each node stores its metadata locally and exchanges it with peers during connection setup. It needs no GPUs, and multiple inference deployments can point at the same cluster. The store nodes reach the master through the `mooncake-master` `Service`. + +```yaml +apiVersion: apps/v1 +kind: Deployment +metadata: + name: mooncake-master + labels: + app: mooncake-master +spec: + replicas: 1 + selector: + matchLabels: + app: mooncake-master + template: + metadata: + labels: + app: mooncake-master + spec: + containers: + - name: mooncake-master + image: lmsysorg/sglang:v0.5.5 + command: ["mooncake_master"] + args: + - --rpc_address + - $(POD_IP) + - --rpc_port + - "50051" + - --metrics_port + - "9003" + env: + - name: POD_IP + valueFrom: + fieldRef: + fieldPath: status.podIP + ports: + - name: rpc + containerPort: 50051 + - name: metrics + containerPort: 9003 + readinessProbe: + tcpSocket: + port: 50051 + initialDelaySeconds: 10 + periodSeconds: 10 +--- +apiVersion: v1 +kind: Service +metadata: + name: mooncake-master + labels: + app: mooncake-master +spec: + type: ClusterIP + selector: + app: mooncake-master + ports: + - name: rpc + port: 50051 + targetPort: 50051 + - name: metrics + port: 9003 + targetPort: 9003 +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: mooncake-store + labels: + app: mooncake-store +spec: + replicas: 3 + selector: + matchLabels: + app: mooncake-store + template: + metadata: + labels: + app: mooncake-store + spec: + containers: + - name: mooncake-store + image: lmsysorg/sglang:v0.5.5 + command: ["python3", "-m", "mooncake.mooncake_store_service"] + args: ["--port", "8088"] + env: + - name: MOONCAKE_MASTER + value: "mooncake-master:50051" + - name: MOONCAKE_TE_META_DATA_SERVER + value: "P2PHANDSHAKE" + - name: MOONCAKE_GLOBAL_SEGMENT_SIZE + value: "10gb" + - name: MOONCAKE_LOCAL_BUFFER_SIZE + value: "0" + - name: MOONCAKE_PROTOCOL + value: "rdma" + resources: + requests: + memory: "16Gi" + limits: + memory: "16Gi" +``` + +See [Notes](#notes) for capacity, protocol (TCP/RDMA), and metadata guidance. + +**Verify:** + +```bash +kubectl get pods -l app=mooncake-master +kubectl get pods -l app=mooncake-store +# Master metrics summary: +kubectl port-forward svc/mooncake-master 9003:9003 & +curl -s http://localhost:9003/metrics/summary +``` + +## Notes + +**Metadata (P2P handshake).** These manifests use Mooncake's P2P handshake (`MOONCAKE_TE_META_DATA_SERVER: P2PHANDSHAKE`): each node stores Transfer Engine metadata locally and exchanges it with peers during connection setup, so there is nothing extra to run and the `mooncake-master` needs no `--*http_metadata_server*` flags. This is the recommended starting point. For large or long-lived clusters, switch the store nodes to the master's embedded HTTP metadata server or an external etcd/Redis instead; see the store guide's [Deployment Scenarios](../mooncake-store-deployment-guide.md#deployment-scenarios). + +**High availability.** A single `mooncake-master` is a single point of failure. For HA, see the store guide's [High Availability](../mooncake-store-deployment-guide.md#deployment-scenarios) section for etcd/Redis backends. + +**TCP vs RDMA.** `MOONCAKE_PROTOCOL` selects the fabric. These manifests use `rdma`. Granting pods RDMA access is cluster-specific and not fully wired into the YAML above; the production reference (see [RBG Integration](rbg-integration)) does it with `hostNetwork: true`, a hostPath mount of `/dev/infiniband`, `privileged` + `IPC_LOCK`/`SYS_RESOURCE`, and an explicit NIC list via `MOONCAKE_DEVICE=`. (A device-plugin `rdma/hca` resource with `MC_MS_AUTO_DISC` / `MC_MS_FILTERS` auto-discovery is an alternative on clusters set up that way.) Switch `MOONCAKE_PROTOCOL` to `tcp` on clusters without an RDMA fabric. + +**Capacity.** Keep `MOONCAKE_GLOBAL_SEGMENT_SIZE` within each pod's memory `limit`. A pure store node issues no `Get`/`Put` itself, so its `MOONCAKE_LOCAL_BUFFER_SIZE` is small; the production RDMA reference sets a modest non-zero buffer (`67108864` = 64 MiB) rather than `0`. + +**Images.** The example uses `lmsysorg/sglang:v0.5.5`. This tag is **not** a reproducible pin — for production, replace it with a verified tag or digest. diff --git a/docs/source/deployment/kubernetes-deployment-guide/rbg-integration.md b/docs/source/deployment/kubernetes-deployment-guide/rbg-integration.md new file mode 100644 index 0000000000..37367df395 --- /dev/null +++ b/docs/source/deployment/kubernetes-deployment-guide/rbg-integration.md @@ -0,0 +1,274 @@ +# RBG Integration + +This page covers the same Store + Transfer Engine P/D scenario as the [main guide](index), deployed with the [sgl-project/rbg](https://github.com/sgl-project/rbg) operator instead of the vanilla `Deployment` / `Service` manifests on the [Mooncake on Kubernetes](mooncake-on-kubernetes) page. Install the RBG operator before applying any `RoleBasedGroup`. + +## RBG example + +The upstream RBG repository ships ready-to-use examples for running Mooncake on RBG: + +- [sgl-pd-disagg-with-mooncake-te.yaml](https://github.com/sgl-project/rbg/blob/main/examples/inference/ecosystem/mooncake/mooncake-transfer-engine/sgl-pd-disagg-with-mooncake-te.yaml) — SGLang P/D disaggregation using Mooncake's Transfer Engine for KV transfer. +- [vllm-pd-disagg-with-mooncake-te.yaml](https://github.com/sgl-project/rbg/blob/main/examples/inference/ecosystem/mooncake/mooncake-transfer-engine/vllm-pd-disagg-with-mooncake-te.yaml) — vLLM P/D disaggregation using Mooncake's Transfer Engine for KV transfer. +- [standalone-mooncake-store.yaml](https://github.com/sgl-project/rbg/blob/main/examples/inference/ecosystem/mooncake/mooncake-store/standalone-mooncake-store.yaml) — the shareable standalone Mooncake Store cluster (`mooncake-master` + `mooncake-store` roles). + +For background on the integration, see the RBG [Mooncake integration KEP](https://github.com/sgl-project/rbg/blob/main/keps/74-mooncake-integration/README.md). + +## Production Mooncake Cluster Example + +A production Mooncake cluster on RBG (`workloads.x-k8s.io/v1alpha2`): a Mooncake **Store** (one `master` plus NUMA-split `store` pods) and SGLang **prefill/decode** engines, all over RDMA. The two RBGs below are the Mooncake-side backend; a router/gateway (out of scope for this page) fronts the prefill/decode endpoints to make the P/D deployment servable — see the [overview](index) and the [Prefill/Decode Disaggregation quick start](../integrations/sglang/hicache-quick-start.md). + +| Group | Roles | Purpose | +|---|---|---| +| Store (`qwen3-0`) | `master`, `store-1000gb` | Mooncake Store — the master coordinator plus NUMA-split store pods contributing the DRAM KV pool | +| Workers (`sglang-workers-0`) | `prefill`, `decode` | SGLang engines: `prefill` is a Store/HiCache client **and** P/D transfer; `decode` does P/D transfer only | + +```{caution} +These manifests are sanitized **excerpts** that show the structure and the Mooncake wiring — they cannot be applied directly. The `decode` role and the second NUMA store container are abbreviated to comments, and image / paths / devices are `<…>` placeholders. Fill them in against your own cluster before applying. +``` + +### 1. Store group — master + NUMA store + +The Store `RoleBasedGroup` has a `master` (the Store coordinator) and a `store-` role whose pods each run two NUMA-pinned store processes. The RBG operator creates a Service `s-qwen3-0-master` for the master role, so clients reach it at `s-qwen3-0-master:50051` — no Service object of your own. Node scheduling uses custom `kvcache.ai/master` and `kvcache.ai/store` labels so a node belongs to exactly one role/size. + +```yaml +apiVersion: workloads.x-k8s.io/v1alpha2 +kind: RoleBasedGroup +metadata: + name: qwen3-0 + namespace: default + labels: { app.kubernetes.io/part-of: mooncake } +spec: + roles: + - name: master + replicas: 1 + standalonePattern: + template: + metadata: + labels: { role: master, app.kubernetes.io/instance: qwen3-0 } + spec: + affinity: + nodeAffinity: + requiredDuringSchedulingIgnoredDuringExecution: + nodeSelectorTerms: + - matchExpressions: + - { key: kvcache.ai/master, operator: In, values: [qwen3_0_master] } + containers: + - name: master + image: + command: + - sh + - -c + - | + ulimit -n 1048576 + ulimit -l unlimited + mooncake_master \ + --rpc_address=$(POD_IP) \ + --rpc_port=50051 \ + --eviction_high_watermark_ratio=0.9 \ + --default_kv_lease_ttl=10000 + env: + - { name: POD_IP, valueFrom: { fieldRef: { fieldPath: status.podIP } } } + - { name: NVIDIA_VISIBLE_DEVICES, value: "void" } # master needs no GPU + securityContext: + # master is an RPC coordinator with no RDMA data path — it does not need + # `privileged`. IPC_LOCK/SYS_RESOURCE cover the `ulimit -l unlimited` / mlock above. + privileged: false + capabilities: { add: ["IPC_LOCK", "SYS_RESOURCE"] } + livenessProbe: + exec: { command: ["/bin/sh", "-c", "pgrep -x mooncake_master >/dev/null"] } + initialDelaySeconds: 20 + periodSeconds: 15 + ports: + - { containerPort: 50051, name: http } + - { containerPort: 9003, name: metrics } + + - name: store-1000gb + replicas: 3 + standalonePattern: + template: + metadata: + labels: + role: store-1000gb + app.kubernetes.io/instance: qwen3-0 + app.kubernetes.io/part-of: mooncake + spec: + affinity: + nodeAffinity: + requiredDuringSchedulingIgnoredDuringExecution: + nodeSelectorTerms: + - matchExpressions: + - { key: kvcache.ai/store, operator: In, values: [qwen3_0_store-1000gb] } + podAntiAffinity: # at most one store pod per node across sizes + requiredDuringSchedulingIgnoredDuringExecution: + - labelSelector: + matchExpressions: + - { key: app.kubernetes.io/part-of, operator: In, values: [mooncake] } + - { key: app.kubernetes.io/instance, operator: In, values: [qwen3-0] } + topologyKey: kubernetes.io/hostname + hostNetwork: true # RDMA + dnsPolicy: ClusterFirstWithHostNet + containers: + - name: store-numa0 + image: + command: + - sh + - -c + - | + ulimit -n 1048576 + ulimit -l unlimited + exec numactl --cpunodebind=0 --membind=0 python3 -m mooncake.mooncake_store_service --port=8099 + env: + - { name: MOONCAKE_LOCAL_HOSTNAME, valueFrom: { fieldRef: { fieldPath: status.podIP } } } + - { name: MOONCAKE_MASTER, value: "s-qwen3-0-master:50051" } + - { name: MOONCAKE_TE_META_DATA_SERVER, value: "P2PHANDSHAKE" } + - { name: MOONCAKE_GLOBAL_SEGMENT_SIZE, value: "1000gb" } # DRAM this store contributes + - { name: MOONCAKE_LOCAL_BUFFER_SIZE, value: "67108864" } + - { name: MOONCAKE_PROTOCOL, value: "rdma" } + - { name: MOONCAKE_DEVICE, value: "" } + - { name: MC_ENABLE_DEST_DEVICE_AFFINITY, value: "1" } + # pin the client metrics HTTP server per container (numa0=9300 / numa1=9301); + # under hostNetwork two processes cannot share 9300. + - { name: MOONCAKE_ENABLE_CLIENT_HTTP_SERVER, value: "true" } + - { name: MOONCAKE_CLIENT_HTTP_PORT, value: "9300" } + ports: [{ containerPort: 8099 }] + startupProbe: + exec: { command: ["sh", "-c", "nc -z 127.0.0.1 8099"] } + periodSeconds: 10 + failureThreshold: 90 + livenessProbe: + exec: { command: ["sh", "-c", "nc -z 127.0.0.1 8099"] } + initialDelaySeconds: 10 + periodSeconds: 10 + securityContext: + # least-privilege RDMA: no `privileged` needed — IPC_LOCK/SYS_RESOURCE plus + # the /dev/infiniband device mount below are enough for the Transfer Engine NICs. + privileged: false + capabilities: { add: ["IPC_LOCK", "SYS_RESOURCE"] } + volumeMounts: + - { mountPath: /dev/infiniband, name: ib } + # store-numa1: identical, but `numactl --cpunodebind=1 --membind=1`, --port=8100, + # containerPort 8100, and MOONCAKE_CLIENT_HTTP_PORT=9301. + volumes: + - { name: ib, hostPath: { path: /dev/infiniband, type: DirectoryOrCreate } } +``` + +### 2. Inference workers — prefill + decode + +`sglang-workers-0` runs the SGLang PD engines, and the two roles are wired **differently**: + +- **`prefill`** is the Store/HiCache client. It sets `MOONCAKE_MASTER=s-qwen3-0-master:50051` (the Store master), `MOONCAKE_TE_META_DATA_SERVER=P2PHANDSHAKE` (the Transfer Engine coordinates P/D directly, no metadata server), `MOONCAKE_PROTOCOL=rdma` + `MOONCAKE_DEVICE`, and `MOONCAKE_GLOBAL_SEGMENT_SIZE=0` (a **pure client** — it contributes no DRAM; the store pods do). It launches with `--enable-hierarchical-cache --hicache-storage-backend mooncake`. The prefill manifest is shown below. +- **`decode`** only participates in the P/D Transfer Engine — `--disaggregation-mode decode --disaggregation-ib-device`. It is **not** a Store/HiCache client: it sets no `MOONCAKE_MASTER` and enables no hierarchical cache. + +```yaml +apiVersion: workloads.x-k8s.io/v1alpha2 +kind: RoleBasedGroup +metadata: + name: sglang-workers-0 + namespace: default +spec: + roles: + - name: prefill + replicas: 4 + standalonePattern: + template: + metadata: + labels: + app: sglang-worker + rbg.workloads.x-k8s.io/group-name: sglang-workers-0 + rbg.workloads.x-k8s.io/role-name: prefill + spec: + hostNetwork: true # RDMA + dnsPolicy: ClusterFirstWithHostNet + nodeSelector: { deployment: sglang_0_prefill } + containers: + - name: sglang-prefill + image: + command: + - bash + - -c + - | + set -e + ulimit -n 1048576; ulimit -l unlimited + python -m sglang.launch_server \ + --model ${MODEL_PATH} --served-model-name Qwen3-0.6B \ + --host 0.0.0.0 --port 8000 \ + --disaggregation-mode prefill \ + --disaggregation-ib-device $IB_DEVICE_LIST \ + --enable-hierarchical-cache --hicache-storage-backend mooncake \ + --tp 8 --page-size 64 --trust-remote-code \ + --enable-metrics --enable-cache-report + # … model/hardware tuning flags omitted (context length, mem fraction, + # NSA backends, EAGLE speculative decoding, KV-cache dtype, etc.) + env: + # --- pod identity --- + - { name: POD_NAME, valueFrom: { fieldRef: { fieldPath: metadata.name } } } + - { name: POD_IP, valueFrom: { fieldRef: { fieldPath: status.podIP } } } + - { name: MOONCAKE_LOCAL_HOSTNAME, valueFrom: { fieldRef: { fieldPath: status.podIP } } } + - { name: SGLANG_HOST_IP, valueFrom: { fieldRef: { fieldPath: status.podIP } } } + # --- model + fabric --- + - { name: MODEL_PATH, value: /models/Qwen3-0.6B } + - { name: IB_DEVICE_LIST, value: "" } + # --- Mooncake wiring --- + - { name: MOONCAKE_TE_META_DATA_SERVER, value: P2PHANDSHAKE } + - { name: MOONCAKE_MASTER, value: "s-qwen3-0-master:50051" } # the Store group's master + - { name: MOONCAKE_PROTOCOL, value: rdma } + - { name: MOONCAKE_DEVICE, value: "" } + - { name: MOONCAKE_GLOBAL_SEGMENT_SIZE, value: "0" } # pure client, contributes no DRAM + - { name: MC_TE_METRIC, value: "true" } + # … SGLANG_* / MC_* performance tuning omitted (heartbeat, timeouts, + # spec-decoding v2, auto-empty-cache, NCCL, JIT, CPU affinity, PRC port range) … + ports: + - { containerPort: 8000, name: http } + - { containerPort: 8998, name: bootstrap } + readinessProbe: + tcpSocket: { port: 8000 } + initialDelaySeconds: 30 + periodSeconds: 10 + resources: + limits: { nvidia.com/gpu: "8" } + requests: { nvidia.com/gpu: "8" } + securityContext: + # least-privilege RDMA: no `privileged` needed — IPC_LOCK/SYS_RESOURCE plus + # the /dev/infiniband device mount below are enough for the NICs. + privileged: false + capabilities: { add: ["IPC_LOCK", "SYS_RESOURCE"] } + volumeMounts: + - { mountPath: /models, name: model } + - { mountPath: /dev/shm, name: dshm } + - { mountPath: /dev/infiniband, name: ib } + volumes: + - { name: model, hostPath: { path: , type: DirectoryOrCreate } } + - { name: dshm, emptyDir: { medium: Memory, sizeLimit: 1300Gi } } + - { name: ib, hostPath: { path: /dev/infiniband, type: DirectoryOrCreate } } + + - name: decode + replicas: 1 + # Abbreviated excerpt — a real decode PodSpec mirrors the prefill container EXCEPT: + # launch: --disaggregation-mode decode --disaggregation-ib-device $IB_DEVICE_LIST + # (NO --enable-hierarchical-cache / --hicache-storage-backend — decode is + # not a Store client), plus DP/EP attention, low-latency deepep, etc. + # env: NO MOONCAKE_MASTER / MOONCAKE_GLOBAL_SEGMENT_SIZE / hierarchical cache; + # keeps POD_NAME / POD_IP / MODEL_PATH / IB_DEVICE_LIST (P/D transfer only) + # sched: nodeSelector deployment: sglang_0_decode ; dshm sizeLimit 15Gi + standalonePattern: + template: { } # fill in a real PodSpec to deploy +``` + +### Mooncake integration points (recap) + +- **Store master** — `mooncake_master --rpc_address=$(POD_IP) --rpc_port=50051 …`; the RBG operator exposes it as `s-qwen3-0-master`, and its clients (the store pods and the **prefill** engine) set `MOONCAKE_MASTER=s-qwen3-0-master:50051`. +- **Only prefill is a Store client** — `prefill` enables HiCache (`--enable-hierarchical-cache --hicache-storage-backend mooncake`) and connects to the master; `decode` participates only in the P/D Transfer Engine and connects to no Store. +- **Transfer Engine uses P2P handshake** — `prefill` sets `MOONCAKE_TE_META_DATA_SERVER=P2PHANDSHAKE`; the TE side-channel coordinates prefill↔decode directly, so there is no HTTP metadata server here (unlike the [main guide](index)). +- **Store vs client segment size** — store pods set `MOONCAKE_GLOBAL_SEGMENT_SIZE=1000gb` (they own the DRAM pool); the prefill client sets `0` (contributes no DRAM, only `Get`/`Put`). +- **RDMA fabric** — `MOONCAKE_PROTOCOL=rdma` + `MOONCAKE_DEVICE=`, `hostNetwork: true`, `/dev/infiniband` mounted, and the `IPC_LOCK`/`SYS_RESOURCE` capabilities (no `privileged` needed — a `/dev/infiniband` hostPath mount plus those two capabilities is the least-privilege way to reach the NICs; an RDMA device plugin is an alternative). +- **NUMA split** — each store pod runs two `mooncake_store_service` processes, each `numactl`-pinned to one NUMA node with its own port and client-metrics port. + +(placeholders)= +### Placeholders + +The example values above (`qwen3-0`, `default`, `store-1000gb`, replica counts) are yours to change. The `<…>` placeholders are: + +| Placeholder | Replace with | +|---|---| +| `` | the container image (`registry/name:tag`) for this role | +| `` | host directory holding the model weights | +| `` | your RDMA NIC list (e.g. `mlx5_0,mlx5_1,…`) | diff --git a/docs/source/deployment/kv-cache-sharing-and-isolation.md b/docs/source/deployment/kv-cache-sharing-and-isolation.md new file mode 100644 index 0000000000..79d172a6c3 --- /dev/null +++ b/docs/source/deployment/kv-cache-sharing-and-isolation.md @@ -0,0 +1,134 @@ +# KV Cache Sharing and Isolation + +## Overview + +Mooncake Store treats object keys as opaque strings. The inference framework constructs those keys and therefore defines which requests and instances may reuse the same KV cache entries. + +A shared Mooncake Store can safely serve multiple models and releases when each cache-compatible group uses a stable, unique namespace. This is useful for: + +- Hosting multiple model families on the same Store cluster. +- Rolling from one model revision to another without mixing their KV cache. +- Running canary, A/B, quantized, or fine-tuned variants side by side. +- Separating cache reuse by request group with `cache_salt`. +- Applying a Mooncake quota to each logical tenant. + +## Choose the Isolation Boundary + +The available isolation mechanisms are complementary: + +| Boundary | Use it for | Configuration | +|----------|------------|---------------| +| Model | Different model families or incompatible model variants | SGLang `--served-model-name`; vLLM derives the name from the model path | +| Deployment or release | Rolling upgrades, canaries, and environments that use the same model name | SGLang `extra_backend_tag`; vLLM `cache_prefix` | +| Request group | Reuse within one application, user group, or other isolation domain | vLLM API `cache_salt` | +| Mooncake tenant | A logical object namespace with its own quota | Store client `tenant_id` and master `--enable_multi_tenants=true` | + +Instances should use identical values only when their KV cache entries are compatible and are intended to be shared. Prefill and decode instances in the same deployment must therefore use the same values. + +## SGLang + +Use an SGLang version that includes [model-aware Mooncake key isolation](https://github.com/sgl-project/sglang/pull/31920), and set a stable, unique `--served-model-name`: + +```bash +python -m sglang.launch_server \ + --model-path Qwen/Qwen3-8B \ + --served-model-name qwen3-8b \ + --enable-hierarchical-cache \ + --hicache-storage-backend mooncake +``` + +All SGLang instances that should share KV cache entries must use the same `--served-model-name`. + +For two deployments of the same served model, add a release-specific `extra_backend_tag`: + +```bash +python -m sglang.launch_server \ + --model-path /models/Qwen3-8B \ + --served-model-name qwen3-8b \ + --enable-hierarchical-cache \ + --hicache-storage-backend mooncake \ + --hicache-storage-backend-extra-config \ + '{"extra_backend_tag":"production-2026-07"}' +``` + +Keep the tag identical across compatible instances in one release, and change it for a release whose KV cache must not be reused. + +## vLLM + +The vLLM `MooncakeStoreConnector` automatically uses the final component of the model passed to `vllm serve` as its model identifier. For example, both `Qwen/Qwen3-8B` and `/models/Qwen3-8B` produce `Qwen3-8B`. + +Set `cache_prefix` when paths with the same final component must remain separate, or when a deployment needs a release-specific namespace: + +```bash +MOONCAKE_CONFIG_PATH=/path/to/mooncake_config.json \ +vllm serve /models/Qwen3-8B \ + --kv-transfer-config '{ + "kv_connector": "MooncakeStoreConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": { + "cache_prefix": "production-2026-07" + } + }' +``` + +All vLLM instances that should share KV cache entries must use the same model and `cache_prefix`. + +### Request-Level Isolation with `cache_salt` + +vLLM also accepts an optional `cache_salt` in each OpenAI-compatible request: + +```json +{ + "model": "qwen3-8b", + "messages": [ + { + "role": "user", + "content": "Summarize this document." + } + ], + "cache_salt": "application-a" +} +``` + +vLLM inserts the salt into the first KV block hash. Because subsequent block hashes depend on the preceding block, the salt affects the Mooncake keys for the request: + +- Requests with the same salt can reuse compatible cached prefixes. +- Requests with different salts use separate cache entries. +- An omitted or empty salt retains the default shared-cache behavior. + +LMCache uses the same isolation-domain convention when it propagates vLLM's `cache_salt` into its cache object identity: one stable salt defines one sharing group. Avoid generating a random salt for every request unless disabling cross-request reuse is intentional. + +## Rolling Model Upgrades + +During a rolling upgrade, the old and new releases may temporarily use the same Mooncake Store. Give each release a distinct deployment namespace: + +| Release | SGLang `extra_backend_tag` | vLLM `cache_prefix` | +|---------|----------------------------|---------------------| +| Current | `production-2026-07` | `production-2026-07` | +| New | `production-2026-08` | `production-2026-08` | + +A typical rollout is: + +1. Keep all current-release instances on the existing namespace. +2. Start the new prefill and decode instances with the new namespace. +3. Warm the new release and gradually shift traffic to it. +4. Retire the old instances after in-flight requests finish. +5. Allow old cache entries to leave the Store through the normal eviction or cleanup lifecycle. + +This approach prevents the new release from reading KV cache generated by the old release, while preserving cache reuse among instances of the same release. Plan for temporarily higher Store capacity because both releases may be warm at the same time. + +Use a new namespace whenever cache compatibility may have changed, including changes to model weights, KV layout, quantization, adapters, or other inference settings that affect generated KV values. + +## Mooncake Tenants + +Mooncake tenant configuration is independent of the framework-level model, release, and request namespaces. When the master is started with `--enable_multi_tenants=true`, the client `tenant_id` selects a tenant-scoped object namespace and the master applies that tenant's quota during admission. + +Use the same `tenant_id` for framework instances that should share one quota and tenant namespace. See [Multi-Tenant Deployment](multi-tenancy) for configuration details. + +## Operational Checklist + +- Use deterministic names that remain stable across restarts. +- Keep model and release fields identical across compatible prefill, decode, and replica instances. +- Change the release namespace before introducing a cache-incompatible deployment. +- Keep `cache_salt` stable within a group that should benefit from prefix reuse. +- Account for duplicate cache occupancy while old and new releases overlap. diff --git a/docs/source/deployment/mooncake-store-deployment-guide.md b/docs/source/deployment/mooncake-store-deployment-guide.md index 39e3ffc9bf..39ed93850a 100644 --- a/docs/source/deployment/mooncake-store-deployment-guide.md +++ b/docs/source/deployment/mooncake-store-deployment-guide.md @@ -1,140 +1,1230 @@ -# Mooncake Store Deployment & Operations Guide - -This page summarizes useful flags, environment variables, and HTTP endpoints to help advanced users tune Mooncake Master and observe metrics. - -## Master Startup Flags (with defaults) - -- RPC Related - - `--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_conn_timeout_seconds` (int, default `0`): RPC idle connection timeout; `0` disables. - - `--rpc_enable_tcp_no_delay` (bool, default `true`): Enable TCP_NODELAY. - -- Metrics - - `--enable_metric_reporting` (bool, default `true`): Periodically log master metrics to INFO. - - `--metrics_port` (int, default `9003`): HTTP port for `/metrics` endpoints. - -- HTTP Metadata Server For Mooncake Transfer Engine - - `--enable_http_metadata_server` (bool, default `false`): Enable embedded HTTP metadata server. - - `--http_metadata_server_host` (str, default `0.0.0.0`): Metadata bind host. - - `--http_metadata_server_port` (int, default `8080`): Metadata TCP port. - -- Allocation Strategy - - `--allocation_strategy` (str, default `random`): Memory allocation strategy for replica placement. Available options: - - `random`: Pure random selection across segments (baseline, fastest). - - `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` (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. - -- High Availability (optional) - - `--enable_ha` (bool, default `false`): Enable HA (requires etcd). - - `--etcd_endpoints` (str, default empty unless HA config): etcd endpoints, semicolon separated. - - `--client_ttl` (int64, default `10` s): Client alive TTL after last ping (HA mode). - - `--cluster_id` (str, default `mooncake_cluster`): Cluster ID for persistence in HA mode. - -- Task Manager (optional) - - `--max_total_finished_tasks` (uint32, default `10000`): Maximum number of finished tasks to keep in memory. When this limit is reached, the oldest finished tasks will be pruned from memory. - - `--max_total_pending_tasks` (uint32, default `10000`): Maximum number of pending tasks that can be queued in memory. When this limit is reached, new task submissions will fail with `TASK_PENDING_LIMIT_EXCEEDED` error. - - `--max_total_processing_tasks` (uint32, default `10000`): Maximum number of tasks that can be processing simultaneously. When this limit is reached, no new tasks will be popped from the pending queue until some processing tasks complete. - - `--max_retry_attempts` (uint32, default `10`): Maximum number of retry attempts for failed tasks. Tasks that fail with `NO_AVAILABLE_HANDLE` error will be retried up to this many times before being marked as failed. - -- DFS Storage (optional) - - `--root_fs_dir` (str, default empty): DFS mount directory for storage backend, used in Multi-layer Storage Support. - - `--global_file_segment_size` (int64, default `int64_max`): Maximum available space for DFS segments. - -- Snapshot / Restore (optional) - - `--enable_snapshot` (bool, default `false`): Enable periodic snapshot of master metadata data (effective when using the `offset` memory allocator). - - `--snapshot_interval_seconds` (uint64, default `600`): Interval in seconds between periodic snapshots of master data. - - `--snapshot_child_timeout_seconds` (uint64, default `300`): Timeout in seconds for each snapshot child process. - - `--snapshot_retention_count` (uint32, default `2`): Number of recent snapshots to keep. Older snapshots beyond this limit will be automatically deleted. - - `--snapshot_backend_type` (str, required when snapshot enabled): Snapshot storage backend type: `local` for local filesystem, `s3` for S3 storage. - - `--snapshot_backup_dir` (str, default empty): Optional local directory for snapshot backup. If empty (default), local backup is disabled. When set, it serves two purposes: (1) during snapshot persistence, data will be saved locally as a fallback if uploading to the backend fails; (2) during restore, downloaded metadata will also be saved to this directory as a local backup. - - `--enable_snapshot_restore` (bool, default `false`): Enable restore from the latest snapshot at master startup. - - **Environment variable** `MOONCAKE_SNAPSHOT_LOCAL_PATH` (**required** when `--snapshot_backend_type=local`): Persistent directory path for local snapshot storage. This variable **must** be set before starting the master; there is no default value. Example: `export MOONCAKE_SNAPSHOT_LOCAL_PATH=/data/mooncake_snapshots`. - - > **Warning: Managed Directory** - > - > The snapshot storage path (`MOONCAKE_SNAPSHOT_LOCAL_PATH` for local backend, or S3 bucket for S3 backend) is a **managed directory** exclusively controlled by the Mooncake snapshot system. **DO NOT store other files or data in this directory.** Old snapshots exceeding `--snapshot_retention_count` will be automatically and permanently deleted during cleanup. Use a dedicated, isolated directory for snapshot storage to avoid accidental data loss. - -Example (enable embedded HTTP metadata and metrics): +# Mooncake Store Deployment & Tuning Guide + +This guide covers minimal deployment, and operational tuning of Mooncake Store. + +## Architecture Overview + + +![architecture](../image/mooncake-store-preview.png) + +**Master Service** (`mooncake_master`): The central coordinator. It manages cluster membership, allocates object storage across client nodes, and enforces eviction/placement policies. Runs as a standalone process. + +**Client Node**: Each node contributes DRAM (and optionally VRAM/SSD) to form the distributed cache pool. Clients communicate with the master over RPC for control operations (`Put`/`Get`/`Remove`), but transfer actual data directly between each other via the Transfer Engine — the master is never in the data path. + +**Metadata Service**: A separate service (etcd, Redis, or HTTP) used by the Transfer Engine for peer discovery and configuration. The master's embedded HTTP metadata server can replace an external etcd/Redis for simple deployments. We also provide a P2P handshake mechanism (`P2PHANDSHAKE`) that enables decentralized metadata management by storing metadata locally on each node, eliminating the need for a centralized service — this is the simplest metadata handshake method and the recommended starting point (see [Quick Start](#quick-start)). + +For a detailed design discussion, see the [Mooncake Store Design](../design/store/mooncake-store.md). + +--- + +## Quick Start + +Deploy a minimal single-node Mooncake Store in three steps. + +### 1. Start the Metadata Service + +This quick start uses **P2P handshake** — the simplest option, with **nothing to start**: each node exchanges and stores Transfer Engine metadata locally during connection setup. You just pass the literal string `P2PHANDSHAKE` as the client's `metadata_server` (step 3). + +For large or long-lived clusters, use the master's embedded HTTP metadata server or an external etcd/Redis instead — see [Deployment Scenarios](#deployment-scenarios). + +### 2. Start the Master Service + +With P2P handshake the master needs no metadata-server flags: + +```bash +mooncake_master +``` + +On success the master logs a single line like: + +``` +Master service started on port 50051, max_threads=4, ... +``` + +The master's default RPC port is `50051`. (To embed an HTTP metadata server instead of using P2P, add `--enable_http_metadata_server=true --http_metadata_server_port=8080`.) + +(start-a-store-client)= +### 3. Start a Store Client + +A client contributes DRAM (and optionally SSD) to the cluster. The simplest way is to embed Mooncake in a Python process and call `store.setup(...)` with `metadata_server="P2PHANDSHAKE"`: + +```python +from mooncake.store import MooncakeDistributedStore + +store = MooncakeDistributedStore() +store.setup( + local_hostname="localhost", + metadata_server="P2PHANDSHAKE", # decentralized; no metadata service + global_segment_size=3200 * 1024 * 1024, # DRAM contributed to the cluster + local_buffer_size=512 * 1024 * 1024, # Transfer Engine buffer + protocol="tcp", + rdma_devices="", # keyword is rdma_devices (not device_name) + master_server_addr="127.0.0.1:50051", # keyword is master_server_addr +) +``` + +There are **three ways** to run a client — programmatic (above), a standalone `mooncake_store_service` process (configured via `MOONCAKE_*`), and the `mooncake_client` real-client RPC process. See [Reference: Client Configuration & Tuning](#reference-client-configuration-tuning) for all three, with full parameter/env tables. + +**What just happened:** + +1. The client registered itself with the master via RPC. +2. The master allocated a 3.2 GB segment on this node and added it to the cluster's memory pool. +3. The client is now ready to serve `Put`/`Get`/`Remove` requests. + +### Run the Stress Benchmark + +Mooncake Store includes sample programs for validating C++ and Python integrations. The [stress benchmark script](gh-file:mooncake-store/tests/stress_cluster_benchmark.py) can be used to verify a two-role prefill/decode setup. + +Configure the script with command-line flags (run with `--help` for the full list): + +- `--local-hostname`: the local machine's reachable IP address or hostname. +- `--metadata-server`: the Transfer Engine metadata service, e.g. `P2PHANDSHAKE`, `http://127.0.0.1:8080/metadata`, or an etcd address. +- `--master-server`: the Mooncake Store master address. Use `IP:Port` in default mode, or `etcd://IP:Port;IP:Port;...;IP:Port` in etcd-backed HA mode. +- `--protocol`: transport, `tcp` / `rdma` / `cxl` / `ascend` (defaults to `rdma`). + +Then start the roles: + +```bash +python3 mooncake-store/tests/stress_cluster_benchmark.py --role prefill +python3 mooncake-store/tests/stress_cluster_benchmark.py --role decode +``` + +For RDMA, topology auto-discovery and NIC filters can be passed through environment variables: + +```bash +MC_MS_AUTO_DISC=1 MC_MS_FILTERS="mlx5_1,mlx5_2" python3 mooncake-store/tests/stress_cluster_benchmark.py --role prefill +MC_MS_AUTO_DISC=1 MC_MS_FILTERS="mlx5_1,mlx5_2" python3 mooncake-store/tests/stress_cluster_benchmark.py --role decode +``` + +The absence of errors indicates successful data transfer. + +### Verify Installed Examples + +For a Python integration check, run `mooncake-store/tests/distributed_object_store_provider.py` after starting the metadata service and `mooncake_master`. + +For a C++ integration check, run `build/mooncake-store/tests/client_integration_test` after building tests and starting the required services. + +### Verify + +```bash +# Health check — master metrics endpoint +curl -s http://localhost:9003/metrics/summary + +# List registered clients +# (exposed through the store's Python API or RPC) +``` + +--- + +## Deployment Scenarios + +### Single-Node (TCP) — Development / Quick Evaluation + +The simplest deployment, as shown in [Quick Start](#quick-start). A single `mooncake_master` orchestrates clients over TCP. Suitable for development, testing, and single-host evaluation. ```bash mooncake_master \ --enable_http_metadata_server=true \ --http_metadata_server_host=0.0.0.0 \ - --http_metadata_server_port=8080 \ - --rpc_thread_num=64 \ - --metrics_port=9003 \ - --enable_metric_reporting=true + --http_metadata_server_port=8080 ``` -Example (use free-ratio-first allocation strategy for better load balancing): +Limitation: the master is a single point of failure. If it crashes, cluster operations pause until it is restored. + +--- + +### High-Availability (etcd) — Production HA + +Runs a cluster of master instances coordinated through etcd. If the leader fails, the remaining instances elect a new leader automatically. ```bash +# Start each master instance with: mooncake_master \ - --allocation_strategy=free_ratio_first \ + --enable_ha=true \ + --ha_backend_type=etcd \ + --ha_backend_connstring="10.0.0.1:2379;10.0.0.2:2379;10.0.0.3:2379" \ + --enable_oplog=true \ + --rpc_address=10.0.0.1 +``` + +Each instance must specify its own reachable `--rpc_address`. `--etcd_endpoints` is still accepted as a backward-compatible alias for the etcd HA backend connection string when `--ha_backend_connstring` is empty. The etcd cluster used for HA can be shared with or separate from the Transfer Engine's metadata etcd. + +**Client addressing:** to reach an HA cluster, clients must use the `etcd://` master-address form (so they can discover the current leader) instead of a single `IP:Port` — set `master_server_addr` (Method A) / `MOONCAKE_MASTER` (Method B) / `--master_server_address` (Method C) to `etcd://10.0.0.1:2379;10.0.0.2:2379;...`. + +--- + +### High-Availability (Redis) — Alternative HA Backend + +Same HA semantics but using Redis instead of etcd for leader election: + +```bash +mooncake_master \ + --enable_ha=true \ + --ha_backend_type=redis \ + --ha_backend_connstring="redis://127.0.0.1:6379" \ + --rpc_address=10.0.0.1 +``` + +**Client addressing:** clients reach a Redis-backed HA cluster with the `redis://connstring` master-address form (e.g. `redis://127.0.0.1:6379`) for `master_server_addr` / `MOONCAKE_MASTER` / `--master_server_address`, instead of a single `IP:Port`. Redis is used only for leader election here. OpLog replication currently requires `ha_backend_type=etcd`. + + +--- + +### Snapshot & Restore — Backup / Disaster Recovery + +```{caution} +Metadata Snapshot And Restore is experimental feature. +``` + +Periodically persist master metadata to local disk or S3, enabling recovery from a recent snapshot after a crash. + +```bash +export MOONCAKE_SNAPSHOT_LOCAL_PATH=/data/mooncake_snapshots + +mooncake_master \ + --enable_snapshot=true \ + --snapshot_interval_seconds=300 \ + --snapshot_retention_count=5 \ + --snapshot_object_store_type=local \ + --enable_snapshot_restore=true +``` + +--- + +### Tiered Storage with SSD Offload — Cost-Effective Capacity + +Extends the cache pool from DRAM to SSD while keeping normal reads and writes on the distributed memory path. With `--enable_offload=true`, completed memory writes are queued for asynchronous SSD persistence through the master control plane. Set `--offload_on_evict=true` to defer that SSD write until the memory eviction path selects an object for reclamation. When `--promotion_on_hit=true`, SSD-only objects can be promoted back to DRAM after repeated reads; admission is gated by `--promotion_admission_threshold`. + +```bash +mooncake_master \ + --enable_offload=true \ + --offload_on_evict=true \ + --promotion_on_hit=true \ + --promotion_admission_threshold=2 \ --enable_http_metadata_server=true \ --http_metadata_server_port=8080 ``` -**Tips:** +Do not set `--root_fs_dir` with `--enable_offload=true`. `--root_fs_dir` is a legacy parameter from an older persistence path and may cause issues on the SSD offload path. Configure each real client's offload directory with `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` instead. + +--- + +### CXL-Aware Allocation — Memory Tiering + +When the host has CXL-attached memory, the master can preferentially allocate new objects on the CXL tier, reserving local DRAM for latency-sensitive operations. + +```bash +mooncake_master \ + --enable_cxl=true \ + --cxl_path=/dev/dax0.0 \ + --cxl_size=17179869184 \ + --allocation_strategy=cxl +``` + +--- -In addition to command-line flags, the Master also supports configuration via JSON and YAML files. For example: +### Container / Dynamic Network Interface + +When the master runs in a container with a dynamic IP, use `--rpc_interface` to resolve the RPC address from a stable interface name: ```bash mooncake_master \ - --config_path=mooncake-store/conf/master.yaml + --rpc_interface=eth0 \ + --enable_http_metadata_server=true \ + --http_metadata_server_host=0.0.0.0 \ + --http_metadata_server_port=8080 ``` -## Metrics Endpoints +The master resolves the current IPv4 address of `eth0` at startup and uses it as the advertised RPC address. + +--- + +## High Availability (HA) + +Mooncake Store supports a Primary-Standby HA model with batch-record OpLog replication. The active Primary serves traffic and writes ordered batches to etcd. Standby nodes poll the durable batch prefix and apply each entry in strict sequence order. + +### HA Architecture + +``` ++------------------+ etcd batch records +---------------+ +| Primary | --------------------------> | Standby | +| OrderedOpLogWriter| durable_prefix | OpLogApplier | +| MasterService | | MetadataStore | ++------------------+ +---------------+ + ^ | + | Leadership Election | + +---------------- etcd/redis/k8s ------------------+ +``` + +### HA Configuration + +HA leadership and metadata replication are configured separately: + +- The HA coordinator elects the active master. Configure it with `--enable_ha`, `--ha_backend_type`, `--ha_backend_connstring`, and `--cluster_id`. For `ha_backend_type=etcd`, legacy `--etcd_endpoints` is used only when `--ha_backend_connstring` is empty. +- The optional batch-record OpLog persists metadata mutations so standby masters can catch up and later be promoted. Enable it explicitly with `--enable_oplog=true`; it is disabled by default and requires `ha_backend_type=etcd` and a build with `STORE_USE_ETCD`. + + +- `--enable_oplog`: Enable the primary OpLog writer and standby reader. Defaults to `false`. +- `--oplog_poll_interval_ms`: Base polling and retry delay for the batch standby, in milliseconds. +- `--oplog_batch_max_entries`: Maximum number of entries admitted to an ordered batch. Defaults to `1024`. +- `--batch_oplog_retry_timeout_sec`: Maximum consecutive retryable batch-standby failure window in seconds (default `180`). -The master exposes Prometheus-style metrics over HTTP on `--metrics_port`: +For snapshot-based standby bootstrap, also configure: -- `GET /metrics` — Prometheus format (`text/plain; version=0.0.4`). -- `GET /metrics/summary` — Human-readable summary. +- `--enable_snapshot_restore` (bool, default `false`): Enable standby to bootstrap from the latest snapshot at startup. +- `--snapshot_object_store_type` (str): Snapshot object store type: `local` or `s3`. +- `--snapshot_catalog_store_type` (str): Snapshot catalog store type: `embedded` (default) or `redis`. -Examples: +### Standby Bootstrap + +When a Standby starts, it follows this sequence: + +1. **Snapshot Bootstrap** (if `enable_snapshot_restore=true`): + - Load the latest snapshot from the configured catalog and object store. + - Rebuild object metadata and segment state from the snapshot baseline. +2. **OpLog Catch-up**: + - Start from the snapshot's `last_included_seq` (or from 1 if no snapshot). + - Poll `durable_prefix`, read batch records up to that boundary, and apply entries in strict sequence order. + +Supported OpLog entry types: +- `PUT_END`: Object write completion +- `REMOVE`: Object removal +- `PUT_REVOKE`: Object revocation +- `SEGMENT_MOUNT`: Segment mount event +- `SEGMENT_UNMOUNT`: Segment unmount event +- `SEGMENT_UPDATE`: Segment update event + +### Promotion and Failover + +When the Primary fails, the Standby is promoted through the following steps: + +1. **Leadership Lease**: The supervisor must acquire and retain the leadership lease before promotion begins. +2. **Final Prefix Read and Catch-up**: The Standby stops its polling loop, reads `durable_prefix` again, and applies all durable batches. A missing prefix is accepted only when the local applied sequence is zero; otherwise promotion fails closed. +3. **Export Context**: The Standby exports its current state as a `PromotionContext`, including: + - `applied_seq_id`: The latest applied OpLog sequence ID. + - `objects`: All object metadata from the in-memory store. + - `segments`: All segment registry entries. +4. **State Restoration**: The new Primary restores and validates the complete `PromotionContext`, populating metadata shards and the segment manager. A context with zero objects and segments still passes through restoration so that unsupported recovery modes cannot bypass validation. +5. **Serving Gate**: The supervisor revalidates leadership and exposes the RPC service only after restoration succeeds. Promotion, restoration, or leadership validation failure leaves `service_ready=false`, keeps data endpoints unavailable, and releases leadership. Failure to release leadership does not make the candidate serviceable. +6. **Invalid Endpoint Filtering**: During restoration, any replica endpoints that correspond to segments no longer in the registry are automatically filtered out from `GetReplicaList` results. + +This fail-closed behavior is intentional. Older versions could log a restoration error and continue serving from empty or partially restored metadata. That behavior was a correctness bug, not a supported availability fallback: the serving state could disagree with the durable OpLog and poison later recovery attempts. Mooncake does not automatically discard snapshots, OpLog records, or metadata after a recovery error. + +### Example: HA Deployment with etcd + +Primary configuration (`primary.yaml`): + +```yaml +enable_ha: true +ha_backend_type: "etcd" +ha_backend_connstring: "etcd-1:2379;etcd-2:2379;etcd-3:2379" +cluster_id: "mooncake_cluster" +enable_oplog: true +oplog_poll_interval_ms: 1000 +oplog_batch_max_entries: 1024 +enable_snapshot: true +snapshot_object_store_type: "local" +snapshot_catalog_store_type: "embedded" +rpc_port: 50051 +``` + +Standby configuration (`standby.yaml`): + +```yaml +enable_ha: true +ha_backend_type: "etcd" +ha_backend_connstring: "etcd-1:2379;etcd-2:2379;etcd-3:2379" +cluster_id: "mooncake_cluster" +enable_oplog: true +oplog_poll_interval_ms: 1000 +oplog_batch_max_entries: 1024 +enable_snapshot_restore: true +snapshot_object_store_type: "local" +snapshot_catalog_store_type: "embedded" +rpc_port: 50052 +``` + +Environment variable for local snapshot storage: ```bash -curl -s http://:9003/metrics -curl -s http://:9003/metrics/summary +export MOONCAKE_SNAPSHOT_LOCAL_PATH=/data/mooncake_snapshots ``` -## Client/Engine Tuning (Env Vars, with defaults) +Start the cluster: -- Topology discovery (Store Client → Transfer Engine) - - `MC_MS_AUTO_DISC` (default `1`): Auto-discover NIC/GPU topology. Set `0` to disable and provide `rdma_devices` manually. - - `MC_MS_FILTERS` (default empty): Optional comma-separated NIC whitelist when auto-discovery is enabled (e.g., `mlx5_0,mlx5_2`). - - If `MC_MS_AUTO_DISC=0`, pass `rdma_devices` (comma-separated) to the Python `setup(...)` call. +```bash +# Start Primary +mooncake_master --config_path=primary.yaml -- Transfer Engine metrics (disabled by default) - - `MC_TE_METRIC` (default `0`/unset): Set to `1` to enable periodic engine metrics logging. **Note:** Not supported when using Transfer Engine TENT. - - `MC_TE_METRIC_INTERVAL_SECONDS` (default `5`): Positive integer seconds between reports (effective only if metrics enabled). +# Start Standby +mooncake_master --config_path=standby.yaml +``` + +### Recovery from Unusable HA State + +First repair temporary backend, configuration, or snapshot-access failures and restart the affected Standby. If the recovery history is confirmed unusable and losing all cached metadata is acceptable, start a new empty cluster explicitly: -- Client metrics (enabled by default) - - `MC_STORE_CLIENT_METRIC` (default `1`): Client-side metrics on by default; set `0` to disable entirely. - - `MC_STORE_CLIENT_METRIC_INTERVAL` (default `0`): Reporting interval in seconds; `0` collects but does not periodically report. +1. Stop every Primary and Standby process that uses the old `cluster_id`. +2. Confirm that losing the old cache metadata and snapshots is acceptable. +3. Change every node to a new, previously unused `cluster_id`. +4. Start the new cluster and allow applications to repopulate the cache. +5. Keep the old namespace for diagnosis, then remove it separately after confirming that no old process can reconnect. -- Local memcpy optimization (Store transfer path) - - `MC_STORE_MEMCPY` (default `0`/false): Set to `1` to prefer local memcpy when source/destination are on the same client. +Using a new `cluster_id` isolates the new cluster from the old OpLog, durable prefix, producer view, and snapshot namespace. Do not delete individual recovery keys or reuse the old `cluster_id` while any old process may still run. There is no automatic reset-on-restore-failure option. -## Set the Log Level for yalantinglibs coro_rpc and coro_http -By default, the log level is set to warning. You can customize it using the following environment variable: +### Resetting a Legacy OpLog Namespace -`export MC_YLT_LOG_LEVEL=info` +The batch-only implementation does not migrate or read older per-entry OpLog data. Reusing a namespace that contains legacy `latest`, numeric entry, or snapshot sidecar keys is rejected. -This sets the log level for yalantinglibs (including coro_rpc and coro_http) to info. +Reset is destructive: -Available log levels: trace, debug, info, warn (or warning), error, and critical. +1. Stop every Primary and Standby process that uses the cluster ID. +2. Confirm that loss of the old metadata and snapshots is acceptable. +3. Delete the complete `/oplog/{cluster_id}` namespace directly with the operator's etcd tooling. +4. Start the cluster with empty state and batch-record OpLog enabled. + +Do not delete individual compatibility keys while any process is running, and do not retain an old snapshot baseline with a nonzero sequence after deleting `durable_prefix`. + +## Metrics Endpoints + +The master exposes Prometheus-style metrics on `--metrics_port`: + +```bash +# Prometheus format +curl -s http://:9003/metrics + +# Human-readable summary +curl -s http://:9003/metrics/summary +``` + +When tenant quota is enabled, `/metrics` also includes per-tenant quota gauges and quota counters: + +- `mooncake_tenant_quota_requested_bytes{tenant_id}` +- `mooncake_tenant_quota_effective_bytes{tenant_id}` +- `mooncake_tenant_quota_charged_bytes{tenant_id}` +- `mooncake_tenant_quota_admission_closed{tenant_id}` +- `mooncake_tenant_quota_over_quota{tenant_id}` +- `mooncake_tenant_quota_explicit_policy{tenant_id}` +- `mooncake_tenant_quota_reject_total{tenant_id,reason}` +- `mooncake_tenant_evict_bytes_total{tenant_id}` +- `mooncake_tenant_quota_allocatable_capacity_bytes` +- `mooncake_tenant_quota_requested_bytes_sum` +- `mooncake_tenant_quota_effective_bytes_sum` + +--- + +## Tenant Quota Management + +:::{toctree} +:maxdepth: 1 + +Multi-Tenant Deployment +::: + +--- ## Quick Tips - 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. +- For detailed SSD offload configuration (storage backends, eviction policies, io_uring), see the [SSD Offload guide](ssd/ssd-offload). +- For NVMe-oF SSD pool configuration see the [NVMe-oF SSD Pool Deployment Guide](ssd/nvmf-ssd-deployment-guide) +- For the experimental HF3FS USRBIO adapter used by descriptor-based DFS replicas, see the [HF3FS USRBIO adapter guide](../getting_started/plugin-usage/3FS-USRBIO-Plugin). +- For detailed monitoring and observation see [Observability](../getting_started/observability) + +:::{toctree} +:maxdepth: 1 +:hidden: + +KV Cache Sharing and Isolation +SSD Storage +HF3FS USRBIO Adapter (Experimental)<../getting_started/plugin-usage/3FS-USRBIO-Plugin> +../getting_started/observability +::: + +--- + +## Reference: Master Startup Flags + +### RPC + +| Flag | Default | Description | +|------|---------|-------------| +| `--rpc_port` | `0` → effective `50051` | RPC listen port. The literal default is `0`, which falls back to the deprecated `--port` (default `50051`) | +| `--rpc_thread_num` | `0` → effective `min(4, CPU cores)` | RPC worker threads. The literal default is `0`, which falls back to the deprecated `--max_threads` → `min(4, CPU cores)` | +| `--rpc_address` | `0.0.0.0` | RPC bind address | +| `--rpc_interface` | empty | Network interface to resolve RPC address at startup (overrides `--rpc_address`) | +| `--rpc_conn_timeout_seconds` | `0` | Idle connection timeout; `0` disables | +| `--rpc_enable_tcp_no_delay` | `true` | Enable TCP_NODELAY | + +### Logging + +The master uses glog. When `--log_dir` is set, all severities are merged into a single journal file in that directory (`mooncake_master.INFO.-